SlowGuess/NeurIPS_2023_OCR huggingface dataset:
so there is an additional source of corruption, which comes from the numerical errors due to the term  $\mathbf{x}_{t_{i + 1}} - \Delta t\mathcal{P}(\mathbf{x}_{t_{i + 1}}) - \mathbf{x}_{t_i}$ . The conditional distribution  $p_{\sigma}(\cdot |\mathbf{x})$  is only approximately Gaussian. Ideally, the effects of numerical errors are dominated by the Gaussian random noise. However, even small errors may accumulate for longer sequences of inference steps. To account for this, we argue that the multi-step loss  $\mathcal{L}_{\mathrm{multi}}(\theta)$  should be used. During training, with the separation of physics update and denoising, the simulation state is first progressed from time  $t_{i + 1}$  to time  $t_i$  using the reverse physics solver. This only yields a perturbed version of the simulation at time  $t_i$  due to numerical inaccuracies. Then a small Gaussian noise is added and, via the denoising network  $s_\theta (\mathbf{x},t)$ , the simulation state is projected back to the distribution  $p_{t_i}$ , which should also resolve the numerical errors. This is iterated, as discussed in section 2 in the main paper, depending on the sliding window size and location.
ResNet We employ a simple ResNet-like architecture, which is used for the score function  $s_{\theta}(\mathbf{x},t)$  and the convolutional neural network baseline (ResNet) for the stochastic heat equation in section 3.2 as well as in section 3.4 again for the score  $s_{\theta}(\mathbf{x},t)$ .
For experiments with periodic boundary conditions, we apply periodic padding with length 16, i.e., if the underlying 2-dimensional data dimensions are  $N \times N$ , the dimensions after the periodic padding are  $(N + 16) \times (N + 16)$ . We implement the periodic padding by tiling the input three times in  $x$ - and  $y$ -direction and then cropping to the correct sizes. The time  $t$  is concatenated as an additional constant channel to the 2-dimensional input data when this architecture represents the score  $s_{\theta}(\mathbf{x},t)$ .
The encoder part of our network begins with a single 2D-convolution encoding layer with 32 filters, kernel size 4, and no activation function. This is followed by four consecutive residual blocks, each consisting of 2D-convolution, LeakyReLU, 2D-convolution, and LeakyReLU. All 2D convolutions have 32 filters with kernel size four and stride 1. The encoder part ends with a single 2D convolution with one filter, kernel size 1, and no activation. Then, in the decoder part, we begin with a transposed 2D convolution, 32 filters, and kernel size 4. Afterward, there are four consecutive residual blocks, analogous to the residual encoder blocks, but with the 2D convolution replaced by a transposed 2D convolution. Finally, there is a final 2D convolution with one filter and kernel size of 5. Parameter counts of this model and other models are given in table 2.
UNet We use the UNet architecture with spatial dropout as described in [Mue+22], Appendix A.1, for the Bayesian neural network baseline of the stochastic heat equation experiment in section 3.2. The dropout rate is set to 0.25. We do not include batch normalization and apply the same periodic padding as done for our ResNet architecture.
FNO The FNO-2D architecture introduced in [Li+21] with  $k_{\mathrm{max},j} = 12$  Fourier modes per channel is used as a baseline for the stochastic heat equation experiment in section 3.2 and the learned physics surrogate model in section 3.4.
Dil-ResNet The Dil-ResNet architecture is described in [Sta+21], Appendix A. This architecture represents the score  $s_{\theta}(\mathbf{x},t)$  in the buoyancy-driven flow with obstacles experiment in section 3.3. We concatenate the constant time channel analogously to the ResNet architecture. Additionally, positional information is added to the network input by encoding the  $x$ -position and  $y$ -position inside the domain in two separate channels.
Architecture	Parameters
ResNet	330 754
UNet	706 035
Dil-ResNet	336 915
FNO	465 377
Table 2: Summary of architectures.
C 1D Toy SDE
For the 1D toy SDE discussed in section 3.1, we consider the SDE given by
$$
d x = - \left[ \lambda_ {1} \cdot \operatorname {s i g n} (x) x ^ {2} \right] d t + \lambda_ {2} d w, \tag {63}
$$
with  $\lambda_{1} = 7$  and  $\lambda_{2} = 0.03$ . The corresponding reverse-time SDE is
$$
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 {64}
$$
Throughout this experiment,  $p_0$  is a categorical distribution, where we draw either 1 or -1 with the same probability. In figure 8, we show trajectories from this SDE simulated with the Euler-Maruyama method. Trajectories start at 1 or -1 and approach 0 as  $t$  increases.
Neural network architecture We employ a neural network  $s_{\theta}(x,t)$  parameterized by  $\theta$  to approximate the score via the 1-step loss, the multi-step loss, implicit score matching [Hyv05, ISM] and sliced score matching with variance reduction [Son+19, SSM-VR]. In all cases, the neural network is a simple multilayer perceptron with elu activations and five hidden layers with 30, 30, 25, 20, and then 10 neurons for the last hidden layer.
We use the Adam optimizer with standard hyperparameters as described in the original paper [KB15]. The learning rate, batch size, and the number of epochs depend on the data set size (100% with 2500 trajectories, 10%, or 1%) and are chosen to ensure convergence of the training loss.
(a)  $\lambda_{2} = 0$
(b)  $\lambda_{2} = 0.03$
Figure 8: Trajectories from SDE (63) with  $\lambda_{2} = 0$  (a) and  $\lambda_{2} = 0.03$  (b).
Training - 1-step loss For the 1-step loss and all data set sizes, we train for 250 epochs with a learning rate of 10e-3 and batch size of 256. In the first phase, we only keep every 5th point of a trajectory and discard the rest. Then, we again train for 250 epochs with the same batch size and a learning rate of 10e-4 but keep all points. Finally, we finetune the network with 750 training epochs and a learning rate of 10e-5.
Training - multi-step loss For the multi-step loss and  $100\%$  of the data set, we first train with the 1-step loss, which resembles a sliding window size of 2. We initially train for 1000 epochs with a batch size of 512 and a learning rate of 10e-3, where we keep only every 5th point on a trajectory and discard the rest. Then, with a decreased learning rate of 10e-4, we begin training with a sliding window size of  $S = 2$  and increment it every 1000 epochs by one until  $S_{\max} = 10$ . In this phase, we train on all points without any removals.
Training - ISM For ISM, we compute the partial derivative  $\partial s_{\theta}(\mathbf{x})_i / \partial \mathbf{x}_i$  using reverse-mode automatic differentiation in JAX (jax.jacrev). For  $100\%$  and  $10\%$  of the data set, we train for 2 000 epochs with a learning rate of 10e-3 and batch size of 10 000. Then we train for an additional 2 000 epochs with a learning rate 10e-4. For  $1\%$ , increase the number of epochs to 20 000.
Training - SSM-VR For sliced score matching with variance reduction [Son+19, SSM-VR], we use the same training setup as for ISM.
(a) ISM learned score.
(b) Multi-step learned score.
(c) ISM reverse-time SDE trajectories.
(d) Multi-step reverse-time SDE trajectories.
(e) ISM probability flow trajectories.
(f) Multi-step probability flow trajectories.
Figure 9: Comparison of Implicit Score Matching (ISM, left) and our proposed training with the multi-step loss (Multi-step, right). Colormap in (a) and (b) truncated to [-75, 75].
Method	Probability flow ODE	Reverse-time SDE
Data set size	Data set size
100%	10%	1%	100%	10%	1%
multi-step	0.97±0.04	0.91±0.05	0.81±0.01	0.99±0.01	0.94±0.02	0.85±0.06
1-step	0.78±0.16	0.44±0.13	0.41±0.13	0.93±0.05	0.71±0.10	0.75±0.10
ISM	0.19±0.05	0.15±0.15	0.01±0.01	0.92±0.05	0.94±0.01	0.52±0.22
SSM-VR	0.17±0.16	0.49±0.24	0.27±0.47	0.88±0.06	0.94±0.06	0.67±0.23
Table 3: Posterior metric  $Q$  for 1000 predicted trajectories averaged over three runs.
Comparison We directly compare the learned score for the reverse-time SDE trajectories and the probability flow trajectories between ISM and the multi-step loss in figure 9 trained on the full data set. The learned score of ISM and the multi-step loss in figure 9a and figure 9b are visually very similar, showing that our method and loss learn the correct score. Overall, after finetuning both ISM and the multi-step loss, the trajectories of the multi-step loss are more accurate compared to ISM. For example, in figure 9e, a trajectory explodes to negative infinity. Also, trajectories from the multi-step loss end in either  $-1$  or  $1$ , while ISM trajectories are attenuated and do not fully reach  $-1$  or  $1$  exactly, particularly for the probability flow ODE.
Results of Table 1 in Main Paper We include the standard deviations of table 1 from the main paper in table 3 above. The posterior metric  $Q$  is very sensitive to the learned score  $s_{\theta}(\mathbf{x},t)$ . Overall, our proposed multi-step loss gives the most consistent and reliable results.
Empirical verification of Theorem 3.1 For the quadratic SDE equation 63, the analytic score is non-trivial, therefore a direct comparison of the learned network  $s_{\theta}(\mathbf{x},t)$  and the true score  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  is difficult. However, we can consider the SDE with affine drift given by
$$
d x = - \lambda x d x + g d W \tag {65}
$$
with  $\lambda = 0.5$  and  $g\equiv 0.04$ . Because this SDE is affine and there are only two starting points  $-1$  and  $1$ , we can write the distribution of states starting in  $x_0$  as a Gaussian with mean  $\mu (t;x_0) = x_0e^{-\lambda t}$  and variance  $\sigma^2 (t;x_0) = \frac{g^2}{2\lambda} (1 - e^{-2\lambda t})$ , see [Øks03]. Then, the score at time  $t$  and position  $x$  conditioned on the starting point  $x_0$  is  $(x - \mu (t;x_0)) / \sigma^2 (t;x_0)$ . See figure 10 for a visualization of the analytic score and a comparison with the learned score. The learned score from the 1-step training matches the analytic score very well in regions where the data density is sufficiently high.
(a) Paths and density
(b) Analytic score field
(c) 1-step training
(d) Multi-step training
Figure 10: 1-step training score matches analytic score. (a) shows some paths sampled from SDE equation (65) and a contour of the density. (b) is the analytic score field, and (c) and (d) are visualizations of the learned score with 1-step and multi-step training. Scores from the multi-step training correspond to more narrow high-density regions. This implies that during inference, trajectories are pulled more strongly to modes of the training data set distribution than for the 1-step training.
D Stochastic Heat Equation
Spectral solver The physics of the 2-dimensional heat equation for  $\mathbf{x} \in \mathbb{R}^{d \times d}$  can be computed analytically. The solver  $\mathcal{P}_h^{\Delta t}(\mathbf{x})$  simulates the systems  $\mathbf{x}$  forward in time by a fixed  $\Delta t$  using the (shifted) Fourier transformation  $\mathcal{F}$ . In particular, we can implement the solver with
$$
\mathcal {P} _ {h} ^ {\Delta t} (\mathbf {x}) = \mathcal {F} ^ {- 1} (A (\Delta t) \circ \mathcal {F} (\mathbf {x})) ， \tag {66}
$$
where  $\circ$  denotes element-wise multiplication and  $A(\Delta t)_{ij} \in \mathbb{R}^{d \times d}$  is a matrix with entries  $A(\Delta t)_{ij} \coloneqq \exp(-\Delta t \cdot \min(i, j, d - i, j - i))$ . The power spectrum of  $\mathbf{x}$  is scaled down by  $A(\Delta t)$ , and higher frequencies are suppressed more than lower frequencies (for  $\Delta t > 0$ ). If noise is added to  $\mathbf{x}$ , then this means that especially higher frequencies are affected. Therefore the inverse transformation (when  $\Delta t > 0$ ) exponentially scales contributions by the noise, causing significant distortions for the reconstruction of  $\mathbf{x}$ .
Spectral loss We consider a spectral error based on the two-dimensional power spectral density. The radially averaged power spectra  $s_1$  and  $s_2$  for two images are computed as the absolute values of the 2D Fourier transform raised to the second power, which are then averaged based on their distance (in pixels) to the center of the shifted power spectrum. We define the spectral error as the weighted difference between the log of the spectral densities
$$
\mathcal {L} \left(s _ {1}, s _ {2}\right) = \sum_ {k} w _ {k} \left| \log \left(s _ {1, k}\right) - \log \left(s _ {2, k}\right) \right| \tag {67}
$$
with a weighting vector  $w \in \mathbb{R}^d$  and  $w_k = 1$  for  $k \leq 10$  and  $w_k = 0$  otherwise.
Training For inference, we consider the linear time discretization  $t_n = n\Delta t$  with  $\Delta t = 0.2/32$  and  $t_{32} = 0.2$ . During training, we sample a random time discretization  $0 \leq t_0 < t_2' < \dots < t_{31}' < t_{32}$  for each batch based on  $t_n$  via  $t_n' \sim \mathcal{U}(t_n - \Delta t/2, t_n + \Delta t/2)$  for  $n = 1, \dots, 31$  and adjust the reverse physics step based on the time difference  $t_i - t_{i-1}$ . In the first training phase, we consider the multi-step loss with a sliding window size of  $S = 6, 8, \dots, 32$  steps, where we increase  $S$  every two epochs. We use Adam to update the weights  $\theta$  with learning rate  $10^{-4}$ . We finetune the network weights for 80 epochs with an initial learning rate of  $10^{-4}$ , which we reduce by a factor of 0.5 every 20 epochs.
$s_{\theta}$  only version For the 1-step loss, this method is similar to Rissanen et al. [RHS22], which proposes a classical diffusion-like model that generates data from the dynamics of the heat equation. Nonetheless, the implementation details and methodology are analogous to the multi-step loss training, except that the reverse physics step  $\hat{\mathcal{P}}^{-1}$  is not explicitly defined but instead learned by the network  $s_{\theta}(\mathbf{x},t)$  together with the score at the same time. We make use of the same ResNet architecture as the default variant. Except for the reverse physics step, the training setup is identical. Although the network  $s_{\theta}(\mathbf{x},t)$  is not trained with any noise for a larger sliding window with the multi-step loss, we add noise to the simulation states for the SDE inference, while there is no noise for the ODE inference.
**Baseline methods** All other baseline methods are trained for 80 epochs using the Adam optimizer algorithm with an initial learning rate of  $10^{-4}$ , which is decreased by a factor of 0.5 every 20 epochs. For the training data, we consider solutions to the heat equation consisting of initial state  $\mathbf{x}_0$  and end state  $\mathbf{x}_T$ .
Test-time distribution shifts We have tested the effects of test-time distribution shifts for the heat equation experiment. We train the score network  $s_{\theta}$  for a specific combination of diffusivity  $\alpha$  and noise  $g$  and vary both parameters for testing. We modify both the simulator and test ground truth for the updated diffusivity and noise. See figure 11. Overall, for small changes of the parameters, there is very little overfitting. Changes in the reconstruction MSE and spectral error can mainly be explained by making the task itself easier or harder to which our network generalizes nicely, e.g., less noise or higher diffusivity leads to smaller reconstruction error.
Additional results We show visualizations for the predictions of different methods in figure 12 and figure 13.
(a) Reverse-time SDE
(b) Probability flow ODE
Figure 11: Test-time distribution shifts for noise and diffusivity. The correction network  $s_{\theta}$  is trained on diffusivity  $\alpha = 1.0$  and noise  $g \equiv 0.1$ . During testing, we vary the diffusivity of the test data set and simulator as well as the noise for inference. Especially the low spectral errors indicate that the network generalizes well to the new behavior of the physics.
E Buoyancy-driven Flow with Obstacles
We use semi-Lagrangian advection for the velocity and MacCormack advection for the hot marker density within a fixed domain  $\Omega \subset [0,1] \times [0,1]$ . The temperature dynamics of the marker field are modeled with a Boussinesq approximation.
Training We train all networks with Adam and learning rate  $10^{-4}$  with batch size 16. We begin training with a sliding window size of  $S = 2$ , which we increase every 30 epochs by 2 until  $S_{\mathrm{max}} = 20$ .
Separate vs. joint updates We compare a joint update of the reverse physics step and corrector function  $s_{\theta}$ , see figure 7b, and a separate update of reverse physics step and corrector function, see figure 7c. An evaluation regarding the reconstruction MSE and perceptual distance is shown in
(b) Input
(c) SMDP - ODE
(d) ResNet
(e) FNO
(f)  $s_{\theta}$  only - ODE
(g)  $s_{\theta}(\mathbf{x},t)$  only - SDE
(i) SMDP -SDE
(h) BNN
Figure 12: Predictions of different methods for the heat equation problem (example 1 of 2).
(a) Ground truth
(f)  $s_\theta$  only - ODE
(b) Input
(c) SMDP - ODE
(g)  $s_{\theta}$  only - SDE
Figure 13: Predictions of different methods for the heat equation problem (example 2 of 2). Neither the BNN nor the "s $_{\theta}$  only" model can produce small-scale structures.
(i) SMDP -SDE
(d) ResNet
(h) BNN
figure 14. Both training and inference variants achieve advantages over "P-1 only" and "sθ only".
approaches. Overall, there are no apparent differences for the ODE inference performance but slight benefits for the SDE inference when separating physics and corrector update.
Limited-memory BFGS We use numerical optimization of the marker and velocity fields at  $t = 0.35$  to match the target smoke and velocity fields at  $t = 0.65$  using limited-memory BFGS [LN89, LBFGS] and the differentiable forward simulation implemented in phiflow [Hol+20]. Our implementation directly uses the LBFGS implementation provided by torch.optim.LBFGS [Pas+19]. As arguments, we use history_size = 10, max_iter = 4 and line_search_fn = strong_wolfe. Otherwise, we leave all other arguments to the default values. We optimize for 10 steps which takes ca. 240 seconds per sample on a single NVIDIA RTX 2070gpu.
Diffusion posterior sampling DPS An additional baseline for this problem is diffusion posterior sampling for general noisy inverse problems [Chu+23, DPS]. As a first step, we pretrain a diffusion model on the data set of marker and velocity fields at  $t = 0.35$ . We use the mask for obstacle positions as an additional conditioning input to the network to which no noise is applied. Our architecture and training closely resemble Denoising Diffusion Probabilistic Models [HJA20, DDPM]. Our network consists of ca. 18.44 million parameters trained for 100k steps and learning rate  $1 \times 10^{-4}$  using cosine annealing with warm restarts ( $T_0 = 50000$ ,  $\eta_{\mathrm{min}} = 5 \times 10^{-6}$ ). The measurement operator  $\mathcal{A}$  is implemented using our differentiable forward simulation. We consider the Gaussian version of DPS, i.e., Algorithm 1 in [Chu+23] with  $N = 1000$ . We fix the step size  $\zeta_i$  at each iteration  $i$  to  $1 / ||y - \mathcal{A}(\hat{\mathbf{x}}_0(\mathbf{x}_i))||$ . For each inference step, we are required to backpropagate gradients through the diffusion model and the forward simulation. Inference for a single sample requires ca. 5000 seconds on a single NVIDIA RTX 2070gpu.
Additional results We provide more detailed visualizations for the buoyancy-driven flow case in figure 16 and figure 17. These again highlight the difficulties of the reverse physics simulator to recover the initial states by itself. Including the learned corrections significantly improves this behavior.
In figure 15, we also show an example of the posterior sampling for the SDE. It becomes apparent that the inferred small-scale structures of the different samples change. However, in contrast to cases like the heat diffusion example, the physics simulation in this scenario leaves only little room for substantial changes in the states.
(a) Reconstruction MSE  $(t = 0.35)$ .
(b) Perceptual distance LPIPS  $(t = 0.35)$ .
Figure 14: Comparison of separate and joint updates averaged over three runs.
Ground truth
Figure 15: Comparison of SMDP-SDE predictions and ground truth for buoyancy-driven flow at  $t = 0.36$ .
SMDP-SDE
(a) Marker field
(b) Velocity field  $(x$  -direction)
Figure 16: Predictions for buoyancy-driven flow with obstacles (example 1 of 2).
(b) Velocity field  $(x$  -direction)
Figure 17: Predictions for buoyancy-driven flow with obstacles (example 2 of 2).
Figure 18: Steps of Langevin dynamics for  $\epsilon = 2\times 10^{-5}$
Figure 19: Steps with Langevin dynamics for  $\epsilon = 2\times 10^{-5}$  and an additional step with  $\Delta t s_{\theta}(\mathbf{x},t)$  which smoothes the images.
F Isotropic turbulence
Training For the physics surrogate model  $\tilde{\mathcal{P}}^{-1}$ , we employ an FNO neural network, see appendix B, with batch size 20. We train the FNO for 500 epochs using Adam optimizer with learning rate  $10^{-3}$ , which we decrease every 100 epochs by a factor of 0.5. We train  $s_{\theta}(\mathbf{x},t)$  with the ResNet architecture, see appendix B, for 250 epochs with learning rate  $10^{-4}$ , decreased every 50 epochs by a factor of 0.5 and batch size 6.
Refinement with Langevin Dynamics Since the trained network  $s_{\theta}(\mathbf{x},t)$  approximates the score  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$ , it can be used for post-processing strategies [WT11; SE19]. We do a fixed point iteration at a single point in time based on Langevin Dynamics via:
$$
\mathbf {x} _ {t} ^ {i + 1} = \mathbf {x} _ {t} ^ {i} + \epsilon \cdot \nabla_ {\mathbf {x}} \log p _ {t} \left(\mathbf {x} _ {t} ^ {i}\right) + \sqrt {2 \epsilon} z _ {t} \tag {68}
$$
for a number of steps  $T$  and  $\epsilon = 2\times 10^{-5}$ , cf. figure 18 and figure 19. For a prior distribution  $\pi_t$ ,  $\mathbf{x}_t^0\sim \pi_t$  and by iterating (68), the distribution of  $\mathbf{x}_t^T$  equals  $p_t$  for  $\epsilon \rightarrow 0$  and  $T\rightarrow \infty$ . There are some theoretical caveats, i.e., a Metropolis-Hastings update needs to be added in (68), and there are additional regularity conditions [SE19].
Additional results We show additional visualizations of the ground truth and reconstructed trajectories in figure 20 and figure 21.
Figure 20: Predictions for isotropic turbulence (example 1 of 2).
Figure 21: Predictions for isotropic turbulence (example 2 of 2).
[Sta+21] K. Stachenfeld, D. B. Fielding, D. Kochkov, M. Cranmer, et al. “Learned Simulators for Turbulence”. In: International Conference on Learning Representations. 2021.
[Vin11] P. Vincent. “A Connection Between Score Matching and Denoising Autoencoders”. In: Neural Comput. 23.7 (2011), pp. 1661-1674.
[WT11] M. Welling and Y. W. Teh. “Bayesian Learning via Stochastic Gradient Langevin Dynamics”. In: Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28 - July 2, 2011. Omnipress, 2011, pp. 681-688.