SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
d x = - \left[ \lambda_ {1} \cdot \operatorname {s i g n} (x) x ^ {2} - \lambda_ {2} ^ {2} \cdot \nabla_ {x} \log p _ {t} (x) \right] d t + \lambda_ {2} d w. \tag {10}
$$
In figure 2a, we show paths from this SDE simulated with the Euler-Maruyama method. The trajectories approach 0 as  $t$  increases. Given the trajectory value at  $t = 10$ , it is no longer possible to infer the origin of the trajectory at  $t = 0$ .
This experiment allows us to use an analytic reverse simulator:  $\hat{\mathcal{P}}^{-1}(x) = \lambda_1 \cdot \mathrm{sign}(x)x^2$ . This is a challenging problem because the reverse physics step increases quadratically with  $x$ , and  $s_\theta$  has to control the reverse process accurately to stay within the training domain, or paths will explode to infinity. We evaluate each model based on how well the predicted trajectories  $\hat{x}_{0:T}$  match the posterior distribution. When drawing  $\hat{x}_T$  randomly from  $[-0.1, 0.1]$ , we should obtain trajectories with  $\hat{x}_0$  being either  $-1$  or  $1$  with the same likelihood. We assign the label  $-1$  or  $1$  if the relative distance of an endpoint is  $< 10\%$  and denote the percentage in each class by  $\rho_{-1}$  and  $\rho_1$ . As some trajectories miss the target, typically  $\rho_{-1} + \rho_1 < 1$ . Hence, we define the posterior metric  $Q$  as twice the minimum of  $\rho_{-1}$  and  $\rho_1$ , i.e.,  $Q := 2 \cdot \min(\rho_{-1}, \rho_1)$  so that values closer to one indicate a better match with the correct posterior distribution.
Training The training data set consists of 2500 simulated trajectories from 0 to  $T$  and  $\Delta t = 0.02$ . Therefore each training trajectory has a length of  $M = 500$ . For the network  $s_{\theta}(x,t)$ , we consider a multilayer perceptron (MLP) and, as a special case, a grid-based discretization (GRID). The latter is not feasible for realistic use cases and high-dimensional data but provides means for an in-depth analysis of different training variants. For GRID, we discretize the domain  $[0,T] \times [-1.25,1.25]$  to obtain a rectangular grid with  $500 \times 250$  cells and linearly interpolate the solution. The cell centers are initialized with 0. We evaluate  $s_{\theta}$  trained via the 1-step and multi-step losses with  $S_{\mathrm{max}} = 10$ . Details of hyperparameters and model architectures are given in appendix C.
Better extrapolation and robustness from multi-step loss See figure 2c for an overview of the differences between the learned score from MLP and GRID and the effects of the multi-step loss. For the 1-step training with MLP, we observe a clear and smooth score field with two tubes that merge to one at  $x = 0$  as  $t$  increases. As a result, the trajectories of the probability flow ODE and reverse-time SDE converge to the correct value. Training via GRID shows that most cells do not get any gradient updates and remain 0. This is caused by a need for more training data in these regions. In addition, the boundary of the trained region is jagged and diffuse. Trajectories traversing these regions can quickly explode. In contrast, the multi-step loss leads to a consistent signal around the center line at  $x = 0$ , effectively preventing exploding trajectories.
Evaluation and comparison with baselines As a baseline for learning the scores, we consider implicit score matching [Hyv05, ISM]. Additionally, we consider sliced score matching with variance reduction [Son+19, SSM-VR] as a variant of ISM. We train all methods with the same network architecture using three different data set sizes. As can be seen in table 1, the 1-step loss,
which is conceptually similar to denoising score matching, compares favorably against ISM and SSM-VR. All methods perform well for the reverse-time SDE, even for very little training data. Using the multi-step loss consistently gives significant improvements at the cost of a slightly increased training time. Our proposed multi-step training performs best or is on par with the baselines for all data set sizes and inference types. Because the posterior metric  $\mathbf{Q}$  is very sensitive to the score where the paths from both starting points intersect, evaluations are slightly noisy.
Method	Probability flow ODE	Reverse-time SDE
Data set size	Data set size
100%	10%	1%	100%	10%	1%
multi-step	0.97	0.91	0.81	0.99	0.94	0.85
1-step	0.78	0.44	0.41	0.93	0.71	0.75
ISM	0.19	0.15	0.01	0.92	0.94	0.52
SSM-VR	0.17	0.49	0.27	0.88	0.94	0.67
Table 1: Posterior metric  $Q$  for 1000 predicted trajectories averaged over three runs. For standard deviations, see table 3 in the appendix.
Comparison with analytic scores We perform further experiments to empirically verify Theorem 3.1 by comparing the learned scores of our method with analytic scores in appendix C.
(a) Inference
Figure 3: Stochastic heat equation overview. While the ODE trajectories provide smooth solutions with the lowest reconstruction MSE, the SDE solutions synthesize high-frequency content, significantly improving spectral error. The "s $_{\theta}$  only" version without the reverse physics step exhibits a significantly larger spectral error. Metrics in (b) are averaged over three runs.
(b) Evaluation
4.2 Stochastic Heat Equation
The heat equation  $\frac{\partial u}{\partial t} = \alpha \Delta u$  plays a fundamental role in many physical systems. For this experiment, we consider the stochastic heat equation, which slightly perturbs the heat diffusion process and includes an additional term  $g(t)\xi$ , where  $\xi$  is space-time white noise, see Pardoux [Par21, Chapter 3.2]. For our experiments, we fix the diffusivity constant to  $\alpha = 1$  and sample initial conditions at  $t = 0$  from Gaussian random fields with  $n = 4$  at resolution  $32\times 32$ . We simulate the heat diffusion with noise from  $t = 0$  until  $t = 0.2$  using the Euler-Maruyama method and a spectral solver  $\mathcal{P}_h$  with a fixed step size  $\Delta t = 6.25\times 10^{-3}$  and  $g\equiv 0.1$ . Given a simulation end state  $\mathbf{x}_T$ , we want to recover a possible initial state  $\mathbf{x}_0$ . In this experiment, the forward solver cannot be used to infer  $\mathbf{x}_0$  directly in a single step or without corrections since high frequencies due to noise are amplified, leading to physically implausible solutions. We implement the reverse physics simulator  $\tilde{P}^{-1}$  by using the forward step of the solver  $\mathcal{P}_h(\mathbf{x})$ , i.e.  $\tilde{\mathcal{P}}^{-1}(\mathbf{x})\approx -\mathcal{P}_h(\mathbf{x})$ .
Training and Baselines Our training data set consists of 2500 initial conditions with their corresponding trajectories and end states at  $t = 0.2$ . We consider a small ResNet-like architecture based on an encoder and decoder part as representation for the score function  $s_{\theta}(\mathbf{x},t)$ . The spectral solver is implemented via differentiable programming in JAX [SC20], see appendix D. As baseline methods, we consider a supervised training of the same ResNet-like architecture as  $s_{\theta}(\mathbf{x},t)$ , a Bayesian neural network (BNN) as well as a Fourier neural operator (FNO) network [Li+21]. We adopt an  $L_{2}$  loss for all these methods, i.e., the training data consists of pairs of initial state  $\mathbf{x}_0$  and end state  $\mathbf{x}_T$ .
Additionally, we consider a variant of our proposed method for which we remove the reverse physics step  $\mathcal{P}^{-1}$  such that the inversion of the dynamics has to be learned entirely by  $s_\theta$ , denoted by " $s_\theta$  only". We do not compare to ISM and SSM-VR in the following as the data dimensions are too high for both methods to train properly, and we did not obtain stable trajectories during inference.
Reconstruction accuracy vs. fitting the data manifold We evaluate our method and the baselines by considering the reconstruction MSE on a test set of 500 initial conditions and end states. For the reconstruction MSE, we simulate the prediction of the network forward in time with the solver  $\mathcal{P}_h$  to obtain a corresponding end state, which we compare to the ground truth via the  $L_{2}$  distance. This metric has the disadvantage that it does not measure how well the prediction matches the training data manifold. I.e., for this case, whether the prediction resembles the properties of the Gaussian random field. For that reason, we additionally compare
the power spectral density of the states as the spectral loss. An evaluation and visualization of the reconstructions are given in figure 3, which shows that our ODE inference performs best regarding the reconstruction MSE. However, its solutions are smooth and do not contain the necessary small-scale structures. This is reflected in a high spectral error. The SDE variant, on the other hand, performs very well in terms of spectral error and yields visually convincing solutions with only a slight increase in the reconstruction MSE. This highlights the role of noise as a source of entropy in the inference process for SMDP SDE, which is essential for synthesizing small-scale structures. Note that there is a natural tradeoff between both metrics, and the ODE and SDE inference perform best for each of the cases while using an identical set of weights.
(a) ODE trajectories
(b) SDE trajectories
Figure 4: Multi-step  $S_{\mathrm{max}}$  vs. reconstruction MSE averaged over 5 runs.
Multi-step loss is crucial for good performance We performed an ablation study on the maximum window size  $S_{\mathrm{max}}$  in figure 4 for the reconstruction MSE. For both ODE and SDE inference, increasing  $S_{\mathrm{max}}$  yields significant improvements at the cost of slightly increased training resources. This also highlights the importance of using a multi-step loss instead of the 1-step loss  $(S_{\mathrm{max}} = 2)$  for inverse problems with poor conditioning.
We perform further experiments regarding test-time distribution shifts when modifying the noise scale and diffusivity, see appendix D, which showcase the robustness of our methods.
4.3 Buoyancy-driven Flow with Obstacles
Next, we test our methodology on a more challenging problem. For this purpose, we consider deterministic simulations of buoyancy-driven flow within a fixed domain  $\Omega \subset [0,1] \times [0,1]$  and randomly placed obstacles. Each simulation runs from time  $t = 0.0$  to  $t = 0.65$  with a step size of  $\Delta t = 0.01$ . SMDP is trained with the objective of reconstructing a plausible initial state given an end state of the marker density and velocity fields at time  $t = 0.65$ , as shown in figure 5a and figure 5b. We place spheres and boxes with varying sizes at different positions within the simulation domain that do not overlap with the marker inflow. For each simulation, we place one to two objects of each category.
Score matching for deterministic systems During training, we add Gaussian noise to each simulation state  $\mathbf{x}_t$  with  $\sigma_t = \sqrt{\Delta t}$ . In this experiment, no stochastic forcing is used to create the data set, i.e.,  $g \equiv 0$ . By adding noise to the simulation states, the 1-step loss still minimizes a score matching objective in this situation, similar to denoising score matching; see appendix A.3 for a derivation. In the situation without stochastic forcing, during inference, our method effectively alternates between the reverse physics step, a small perturbation, and the correction by  $s_\theta(\mathbf{x}, t)$  which projects the perturbed simulation state back to the distribution  $p_t$ . We find that for the SDE trajectories,  $C = 2$  slightly overshoots, and  $C = 1$  gives an improved performance. In this setting, the "s_\theta only" version of our method closely resembles a denoiser that learns additional physics dynamics.
Training and comparison Our training data set consists of 250 simulations with corresponding trajectories generated with phiflow [Hol+20]. Our neural network architecture for  $s_{\theta}(\mathbf{x},t)$  uses dilated convolutions [Sta+21], see appendix E for details. The reverse physics step  $\hat{\mathcal{P}}^{-1}$  is implemented directly in the solver by using a negative step size  $-\Delta t$  for time integration. For training, we consider the multi-step formulation with  $S_{\mathrm{max}} = 20$ . We additionally compare with solutions from directly optimizing the initial smoke and velocity states at  $t = 0.35$  using the differentiable forward simulation and limited-memory BFGS [LN89, LBFGS]. Moreover, we compare with solutions obtained from diffusion posterior sampling for general noisy inverse problems [Chu+23, DPS] with a pretrained diffusion model on simulation states at  $t = 0.35$ . For the evaluation, we consider a reconstruction MSE analogous to section 4.2 and the perceptual similarity metric LPIPS. The test set contains five simulations. The SDE version yields good results for this experiment but is most likely constrained in performance by the approximate reverse physics step and large step sizes. However, the ODE version outperforms directly inverting the simulation numerically  $(\hat{\mathcal{P}}^{-1}$  only), and when training without the reverse physics step  $(s_{\theta}$  only), as shown in 5c.
4.4 Navier-Stokes with Unknown Dynamics
As a fourth experiment, we aim to learn the time evolution of isotropic, forced turbulence with a similar setup as Li et al. [Li+21]. The training data set consists of vorticity fields from 1000 simulation trajectories from  $t = 0$  until  $T = 10$  with  $\Delta t = 1$ , a spatial resolution of  $64 \times 64$  and
Figure 5: Buoyancy flow case. Ground truth shows the marker density and velocity field in the  $x$ -direction at different points of the simulation trajectory from the test set (a, b). We show reconstructions given the simulation end state at  $t = 0.65$  and provide an evaluation of the reconstructed trajectories based on perceptual similarity (LPIPS) and the reconstruction MSE for three runs (c).
Figure 6: Turbulence case. Comparison of reconstructed trajectories (a) and evaluation of MSE and spectral error for different training variants (b). Our proposed ODE and SDE inference outperforms the learned surrogate model  $\tilde{\mathcal{P}}^{-1}$ . Metrics are averaged over three runs.
Learned reverse step  $\tilde{P}^{-1}$
Simultaneous
■ training of
P-1 and s
Pretrained  $\tilde{p}^{-1}$
SMDP ODE
Pretrained  $\tilde{p}^{-1}$
SMDP SDE 
(b) Evaluation
viscosity fixed at  $\nu = 10^{-5}$ . As before, our objective is to predict a trajectory  $\hat{\mathbf{x}}_{0:M}$  that reconstructs the true trajectory given an end state  $\mathbf{x}_M$ . In this experiment, we pretrain a surrogate for the reverse physics step  $\tilde{\mathcal{P}}^{-1}$  by employing the FNO architecture from [Li+21] trained on the reverse simulation. For pretraining  $\tilde{\mathcal{P}}^{-1}$  we use our proposed training setup with the multi-step loss and  $S_{\mathrm{max}} = 10$  but freeze the score to  $s_\theta (\mathbf{x},t)\equiv 0$ . Then, we train the time-dependent score  $s_\theta (\mathbf{x},t)$  while freezing the reverse physics step. This approach guarantees that any time-independent physics are captured by  $\tilde{\mathcal{P}}^{-1}$  and  $s_\theta (\mathbf{x},t)$  can focus on learning small improvements to  $\tilde{\mathcal{P}}^{-1}$  as well as respond to possibly time-dependent data biases. We give additional training details in appendix F.
Evaluation and training variants For evaluation, we consider the MSE and spectral error of the reconstructed initial state  $\hat{\mathbf{x}}_0$  compared to the reference  $\mathbf{x}_0$ . As baselines, during inference, we employ only the learned surrogate model  $\tilde{\mathcal{P}}^{-1}$  without  $s_\theta$ . In addition to that, we evaluate a variant for which we train both the surrogate model and  $s_\theta(\mathbf{x}, t)$  at the same time. As the two components resemble the drift and score of the reverse-time SDE, this approach is similar to DiffFlow [ZC21], which learns both components in the context of generative modeling. We label this approach simultaneous training. Results are shown in figure 6. Similar to the stochastic heat equation results in section 4.2, the SDE achieves the best spectral error, while the ODE obtains the best MSE. Our proposed method outperforms both the surrogate model and the simultaneous training of the two components.
We presented a combination of learned corrections training and diffusion models in the context of physical simulations and differentiable physics for solving inverse physics problems. We showed its competitiveness, accuracy, and long-term stability in challenging and versatile experiments and motivated our design choices. We considered two variants with complementary benefits for inference: while the ODE variants achieve the best MSE, the SDE variants allow for sampling the posterior and yield an improved coverage of the target data manifold. Additionally, we provided theoretical insights that the 1-step is mathematically equivalent to optimizing the score matching objective. We showed that its multi-step extension maximizes a variational lower bound for maximum likelihood training.
Despite the promising initial results, our work has limitations that open up exciting directions for future work: Among others, it requires simulating the system backward in time step by step, which can be costly and alleviated by reduced order methods. Additionally, we assume that  $\Delta t$  is sufficiently small and the reverse physics simulator is accurate enough. Determining a good balance between accurate solutions with few time steps and diverse solutions with many time steps represents an important area for future research.
Acknowledgements BH, SV, and NT acknowledge funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (LEDA: grant agreement No 758853, and SpaTe: grant agreement No 863850). SV thanks the Max Planck Society for support through a Max Planck Lise Meitner Group. This research was partly carried out on the High Performance Computing resources of the FREYA cluster at the Max Planck Computing and Data Facility (MPCDF) in Garching operated by the Max Planck Society (MPG).
[Um+20] K. Um, R. Brand, Y. (Fei, P. Holl, and N. Thuerey. "Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers". In: Advances in Neural Information Processing Systems. 2020. URL: https://proceedings.neurips.cc/paper/2020/bitstream/43e4e6a6f341e00671e123714de019a8-AAbstract.html.
[Vin11] P. Vincent. “A Connection Between Score Matching and Denoising Autoencoders”. In: Neural Comput. 23.7 (2011), pp. 1661–1674.
[YSM22] R. Yang, P. Srivastava, and S. Mandt. "Diffusion Probabilistic Modeling for Video Generation". In: CoRR abs/2203.09481 (2022). DOI: 10.48550/arXiv.2203.09481. URL: https://doi.org/10.48550/arXiv.2203.09481.
[QC21] Q. Zhang and Y. Chen. "Diffusion normalizing flow". In: Advances in Neural Information Processing Systems 34 (2021), pp. 16280-16291.
[ZDG96] K. Zhou, J. Doyle, and K. Glover. "Robust and optimal control". In: Control Engineering Practice 4.8 (1996), pp. 1189-1190.
A Proofs and Training Methodology
Below we summarize the problem formulation from the main paper and provide details about the training procedure and the derivation of our methodology.
Problem setting Let  $(\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 {11}
$$
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$ .
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})$ .
Assumptions Throughout this paper, we make some additional assumptions to ensure the existence of a unique solution to the SDE (11) and the strong convergence of the Euler-Maruyama method. In particular:
- Finite variance of samples from  $p_0$ :  $\mathbb{E}_{\mathbf{x}_0 \sim p_0} [||\mathbf{x}_0||_2^2] < \infty$
- Lipschitz continuity of  $\mathcal{P}$ :  $\exists K_1 > 0\forall \mathbf{x},\mathbf{y}\in \mathbb{R}^D:||\mathcal{P}(\mathbf{x}) - \mathcal{P}(\mathbf{y})||_2\leq K_1||\mathbf{x} - \mathbf{y}||_2$
- Lipschitz continuity of  $g$ :  $\exists K_2 > 0 \forall t, s \in [0, T] : |g(t) - g(s)| \leq K_3|t - s|$
- Linear growth condition:  $\exists K_{3} > 0\forall \mathbf{x}\in \mathbb{R}^{D}:||\mathcal{P}(\mathbf{x})||_{2}\leq K_{3}(1 + ||\mathbf{x}||_{2})$
-  $g$  is bounded:  $\exists K_4 > 0 \forall t \in [0, T] : |g(t)| \leq K_4$
Euler-Maruyama Method Using Euler-Maruyama steps, we can simulate paths from SDE (11) similar to ordinary differential equations (ODE). Given an initial state  $\mathbf{X}_{t_0}$ , we let  $\hat{\mathbf{X}}_{t_0}^{\Delta t} = \mathbf{X}_{t_0}$  and define recursively
$$
\hat {\mathbf {X}} _ {t _ {m + 1}} ^ {\Delta t} \leftarrow \hat {\mathbf {X}} _ {t _ {m}} ^ {\Delta t} + \Delta t \mathcal {P} \left(\hat {\mathbf {X}} _ {t _ {m}} ^ {\Delta t}\right) + \sqrt {\Delta t} g \left(t _ {m}\right) z _ {t _ {m}}, \tag {12}
$$
where  $z_{t_m}$  are i.i.d. with  $z_{t_m} \sim \mathcal{N}(0, I)$ . For  $t_i \leq t < t_{i+1}$ , we define the piecewise constant solution of the Euler-Maruyama Method as  $\hat{\mathbf{X}}_t^{\Delta t} \coloneqq \hat{\mathbf{X}}_{t_i}^{\Delta t}$ . Let  $\mathbf{X}_t$  denote the solution of the SDE (11). Then the Euler-Maruyama solution  $\hat{\mathbf{X}}_t^{\Delta t}$  converges strongly to  $\mathbf{X}_t$ .
Lemma A.1. [Strong convergence of Euler-Maruyama method] Consider the piecewise constant solution  $\hat{X}_t^{\Delta t}$  of the Euler-Maruyama method. There is a constant  $C$  such that
$$
\sup  _ {0 \leq t \leq T} \mathbb {E} [ | | X _ {t} - \hat {X} _ {t} ^ {\Delta t} | | _ {2} ] \leq C \sqrt {\Delta t}. \tag {13}
$$
Proof. See Kloeden et al. [Klo+92, p. 10.2.2]
A.1 1-step Loss and Score Matching Objective
Theorem A.2. Consider a data set  $\{\mathbf{x}_{0:m}^{(n)}\}_{n=1}^{N}$  with trajectories sampled from SDE (11). Then the 1-step loss
$$
\mathcal {L} _ {\text {s i n g l e}} (\theta) := \frac {1}{M} \mathbb {E} _ {\mathbf {x} _ {0: M}} \left[ \right. \sum_ {m = 0} ^ {M - 1} \left[ \right.\left| \right.\left| \mathbf {x} _ {m} - \left\{\mathbf {x} _ {m + 1} + \Delta t \left[ - \mathcal {P} \left(\mathbf {x} _ {m + 1}\right) + s _ {\theta} \left(\mathbf {x} _ {m + 1}, t _ {m + 1}\right)\right]\right\}\right| ^ {2} \left. \right]\left. \right] \tag {14}
$$
is equivalent to minimizing the score matching objective
$$
\mathcal {J} (\theta) := \int_ {0} ^ {T} \mathbb {E} _ {\mathbf {x} _ {t}} [ \| \nabla_ {\mathbf {x}} \log p _ {t} (\mathbf {x}) - \tilde {s} _ {\theta} (\mathbf {x}, t) \| _ {2} ^ {2} ] d t, \tag {15}
$$
where  $\tilde{s}_{\theta}(\mathbf{x},t) = s_{\theta}(\mathbf{x},t) / g^{2}(t)$  as  $\Delta t\to 0$