SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$\nLeftrightarrow$  : Consider  $\theta^{*}$  such that  $\tilde{s}_{\theta^{*}}(\mathbf{x},t)\equiv \nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  , which minimizes the score matching objective  $\mathcal{J}(\theta)$  . Then fix a time step  $t$  and sample  $\mathbf{x}_t$  and  $\mathbf{x}_{t + \Delta t}$  from the data set. The probability flow solution  $\mathbf{x}_t^p$  based on equation (14) is
$$
\mathbf {x} _ {t} ^ {p} := \mathbf {x} _ {t + \Delta t} + \Delta t \left[ - \mathcal {P} \left(\mathbf {x} _ {t + \Delta t}\right) + s _ {\theta} \left(\mathbf {x} _ {t + \Delta t}, t + \Delta t\right) \right]. \tag {16}
$$
At the same time, we know that the transformation of marginal likelihoods from  $p_{t + \Delta t}$  to  $p_t$  follows the reverse-time SDE [And82]
$$
d \mathbf {x} = \left[ \mathcal {P} (\mathbf {x}) + g ^ {2} (t) \nabla_ {\mathbf {x}} \log p _ {t} (\mathbf {x}) \right] d t + g (t) d W, \tag {17}
$$
which runs backward in time from  $T$  to 0. Denote by  $\hat{\mathbf{x}}_t^{\Delta t}$  the solution of the Euler-Maruyama method at time  $t$  initialized with  $\mathbf{x}_{t + \Delta t}$  at time  $t + \Delta t$ .
Using the triangle inequality for squared norms, we can write
$$
\lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[ | | \mathbf {x} _ {t} - \mathbf {x} _ {t} ^ {p} | | _ {2} ^ {2} \right] \leq 2 \lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[ | | \mathbf {x} _ {t} - \hat {\mathbf {x}} _ {t} ^ {\Delta t} | | _ {2} ^ {2} \right] + 2 \lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[ | | \hat {\mathbf {x}} _ {t} ^ {\Delta t} - \mathbf {x} _ {t} ^ {p} | | _ {2} ^ {2} \right]. \tag {18}
$$
Because of the strong convergence of the Euler-Maruyama method, we have that for the first term of the bound in equation (18)
$$
\lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[\left|\left| \mathbf {x} _ {t} - \hat {\mathbf {x}} _ {t} ^ {\Delta t} \right|\right| _ {2} ^ {2} \right] = 0 \tag {19}
$$
independent of  $\theta$ . At the same time, for the Euler-Maruyama discretization, we can write
$$
\begin{array}{l} \dot {\mathbf {x}} _ {t} ^ {\Delta t} = \mathbf {x} _ {t + \Delta t} + \Delta t \left[ - \mathcal {P} \left(\mathbf {x} _ {t + \Delta t}\right) + g ^ {2} (t + \Delta t) \nabla_ {\mathbf {x}} \log p _ {t + \Delta t} \left(\mathbf {x} _ {t + \Delta t}\right) \right] (20) \\ + g (t + \Delta t) \sqrt {\Delta t} z _ {t + \Delta t}, (21) \\ \end{array}
$$
where  $z_{t + \Delta t}$  is a standard Gaussian distribution, i.e.,  $z_{t + \Delta t} \sim \mathcal{N}(0, I)$ . Therefore, we can simplify the second term of the bound in equation (18)
$$
\begin{array}{l} \lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[\left|\left| \hat {\mathbf {x}} _ {t} ^ {\Delta t} - \mathbf {x} _ {t} ^ {p} \right|\right| _ {2} ^ {2} \right] (22) \\ = \lim  _ {\Delta t \rightarrow 0} \mathbb {E} _ {\mathbf {x} _ {t + \Delta t} \sim p _ {t + \Delta t}, z \sim \mathcal {N} (0, I)} \left[\left|\right| \Delta t g (t + \Delta t) ^ {2} \left[ \nabla_ {\mathbf {x}} \log p _ {t + \Delta t} (\mathbf {x} _ {t + \Delta t}) \right.\right. (23) \\ \left. \left. - \tilde {s} _ {\theta} \left(\mathbf {x} _ {t + \Delta t}, t + \Delta t\right) \right] + g (t + \Delta t) \sqrt {\Delta t} z \right| _ {2} ^ {2} \Bigg ]. (24) \\ \end{array}
$$
If  $\theta^{*}$  minimizes the score matching objective, then  $\tilde{s}_{\theta^{*}}(\mathbf{x},t)\equiv \nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  , and therefore the above is the same as
$$
\lim  _ {\Delta t \rightarrow 0} \mathbb {E} _ {z} [ | | \sqrt {\Delta t} g (t + \Delta t) z | | _ {2} ^ {2} ] = 0. \tag {25}
$$
Combining equations (18), (19) and (25) yields
$$
\lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[\left|\left| \mathbf {x} _ {t} - \mathbf {x} _ {t} ^ {p} \right|\right| _ {2} ^ {2} \right] = 0. \tag {26}
$$
Additionally, since  $g$  is bounded, we even have
$$
\mathbb {E} [ \| \sqrt {\Delta t} g (t + \Delta t) z \| _ {2} ^ {2} ] \leq \mathbb {E} [ \| \sqrt {\Delta t} K _ {4} z \| _ {2} ^ {2} ] = \mathbb {E} [ \| K _ {4} z \| _ {2} ^ {2} ] \Delta t. \tag {27}
$$
For  $\mathcal{L}_{\mathrm{single}}(\theta)$ , using the above bound (27) and strong convergence of the Euler-Maruyama method, we can therefore derive the following bound
$$
\begin{array}{l} \mathcal {L} _ {\text {s i n g l e}} (\theta) = \frac {1}{M} \mathbb {E} \left[ \sum_ {m = 0} ^ {M - 1} \left[ \left| \left| \mathbf {x} _ {m} + \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} \right] \right] (28) \\ \leq \frac {1}{M} M 2 \left[ C \sqrt {\Delta t} + \mathbb {E} _ {z} [ | | K _ {4} z | | _ {2} ^ {2} ] \Delta t \right], (29) \\ \end{array}
$$
which implies that  $\mathcal{L}_{\mathrm{single}}(\theta)\to 0$  as  $\Delta t\rightarrow 0$
"  $\Rightarrow$  " : With the definitions from  $\nVdash$  ", let  $\theta^{*}$  denote a minimizer such that  $\mathcal{L}_{\mathrm{single}}(\theta)\to 0$  as  $\Delta t\rightarrow 0$ , i.e., we assume there is a sequence  $\Delta t_1,\Delta t_2,\ldots$  with  $\lim_{n\to \infty}\Delta t_n = 0$  and a sequence  $\theta_{1},\theta_{2},\dots$ , where  $\theta_{n}$  is a global minimizer to the objective  $\mathcal{L}_{\mathrm{single}}^{\Delta t_n}(\theta)$  that depends on the step size  $\Delta t_{n}$ . If there is  $\theta^{*}$  such that  $s_{\theta^{*}}(\mathbf{x},t)\equiv \nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ , then  $\mathcal{L}_{\mathrm{single}}^{\Delta t_n}(\theta_n)\leq \mathcal{L}_{\mathrm{single}}^{\Delta t_n}(\theta^*)$ . From  $\nVdash$  " we know that  $\lim_{n\to \infty}\mathcal{L}_{\mathrm{single}}^{\Delta t_n}(\theta^*) = 0$  and therefore  $\lim_{n\to \infty}\mathcal{L}_{\mathrm{single}}^{\Delta t_n}(\theta_n) = 0$ . This implies that each summand of  $\mathcal{L}_{\mathrm{single}}(\theta)$  also goes to zero as  $\Delta t\to 0$ , i.e.,  $\lim_{\Delta t\to 0}\mathbb{E}\left[||\mathbf{x}_t - \mathbf{x}_t^p ||_2^2\right] = 0$ . Again, with the triangle inequality for squared norms, we have that
$$
\lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[ | | \hat {\mathbf {x}} _ {t} ^ {\Delta t} - \mathbf {x} _ {t} ^ {p} | | _ {2} ^ {2} \right] \leq 2 \lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[ | | \mathbf {x} _ {t} - \hat {\mathbf {x}} _ {t} ^ {\Delta t} | | _ {2} ^ {2} \right] + 2 \lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[ | | \mathbf {x} _ {t} - \mathbf {x} _ {t} ^ {p} | | _ {2} ^ {2} \right]. \tag {30}
$$
By the strong convergence of the Euler-Maruyama method and  $\theta = \theta^{*}$ , we obtain
$$
\lim  _ {\Delta t \rightarrow 0} \mathbb {E} \left[\left|\left| \hat {\mathbf {x}} _ {t} ^ {\Delta t} - \mathbf {x} _ {t} ^ {p} \right|\right| _ {2} ^ {2} \right] = 0. \tag {31}
$$
At the same time, for fixed  $\Delta t > 0$ , we can compute
$$
\begin{array}{l} \mathbb {E} \left[ \left| \left| \hat {\mathbf {x}} _ {t} ^ {\Delta t} - \mathbf {x} _ {t} ^ {p} \right| \right| _ {2} ^ {2} \right] (32) \\ = \mathbb {E} _ {\mathbf {x} _ {t + \Delta t}, z \sim \mathcal {N} (0, I)} [ | | \Delta t g (t + \Delta t) ^ {2} [ \nabla_ {\mathbf {x}} \log p _ {t + \Delta t} (\mathbf {x} _ {t + \Delta t}) (33) \\ \left. - s _ {\theta} \left(\mathbf {x} _ {t + \Delta t}, t + \Delta t) \right] + \sqrt {\Delta t} g (t + \Delta t) z \| _ {2} ^ {2} \right] (34) \\ = \Delta t g (t + \Delta t) \mathbb {E} _ {\mathbf {x} _ {t + \Delta t}, z \sim \mathcal {N} (0, I)} [ \| \Delta t ^ {3 / 2} g (t + \Delta t) ^ {5 / 2} \left[ \nabla_ {\mathbf {x}} \log p _ {t + \Delta t} (\mathbf {x} _ {t + \Delta t}) \right. (35) \\ \left. - s _ {\theta} \left(\mathbf {x} _ {t + \Delta t}, t + \Delta t\right) \right] + z \left| \right| _ {2} ^ {2} ]. (36) \\ \end{array}
$$
For fixed  $\mathbf{x}_{t + \Delta t}$ , the distribution over  $z\sim \mathcal{N}(0,I)$  in equation (36) correspond to a noncentral chi-squared distribution [JKB95, Chapter 13.4], whose mean can be calculated as
$$
\begin{array}{l} \mathbb {E} _ {z \sim \mathcal {N} (0, I)} \left[ \left| \left| \Delta t ^ {3 / 2} g (t + \Delta t) ^ {5 / 2} \left[ \nabla_ {\mathbf {x}} \log p _ {t + \Delta t} \left(\mathbf {x} _ {t + \Delta t}\right) - s _ {\theta} \left(\mathbf {x} _ {t + \Delta t}, t + \Delta t\right) \right] + z \right| \right| _ {2} ^ {2} \right] (37) \\ = \left| \left| \Delta t ^ {3 / 2} g (t + \Delta t) ^ {5 / 2} \left[ \nabla_ {\mathbf {x}} \log p _ {t + \Delta t} \left(\mathbf {x} _ {t + \Delta t}\right) - s _ {\theta} \left(\mathbf {x} _ {t + \Delta t}, t + \Delta t\right) \right] \right| \right| _ {2} ^ {2} + D. (38) \\ \end{array}
$$
For each  $\Delta t > 0$ , the above is minimized if and only if  $\nabla_{\mathbf{x}}\log p_{t + \Delta t}(\mathbf{x}_{t + \Delta t}) = s_{\theta}(\mathbf{x}_{t + \Delta t},t + \Delta t)$ .
A.2 Multi-step Loss and Maximum Likelihood Training
We now extend the 1-step formulation from above to multiple steps and discuss its relation to maximum likelihood training. For this, we consider our proposed probability flow ODE defined by
$$
d \mathbf {x} = \left[ \mathcal {P} (\mathbf {x}) + s _ {\theta} (\mathbf {x}, t) \right] d t \tag {39}
$$
and for  $t_i < t_j$  define  $p_{t_i}^{t_j,\mathrm{ODE}}$  as the distribution obtained by sampling  $\mathbf{x} \sim p_{t_j}$  and integrating the probability flow with network  $s_\theta(\mathbf{x}, t)$  equation (7) backward in time until  $t_i$ . We can choose two arbitrary time points  $t_i$  and  $t_j$  with  $t_i < t_j$  because we do not require fixed start and end times of the simulation.
The maximum likelihood training objective of the probability flow ODE (7) can be written as maximizing
$$
\mathbb {E} _ {\mathbf {x} _ {t _ {i}} \sim p _ {t _ {i}}} \left[ \log p _ {t _ {i}} ^ {p _ {t _ {j}}, \text {O D E}} \left(\mathbf {x} _ {t _ {i}}\right) \right]. \tag {40}
$$
Our proposed multi-step loss is based on the sliding window size  $S$ , which is the length of the sub-trajectory that we aim to reconstruct with the probability flow ODE (7). The multi-step loss is defined as
$$
\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 {41}
$$
where we compute the expectation by sampling  $\mathbf{x}_{0:m}$  from the training data set and  $\hat{\mathbf{x}}_{i:i + S - 1}$  is the predicted sub-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[ - \mathcal {P} \left(\hat {\mathbf {x}} _ {i + S - j}\right) + s _ {\theta} \left(\hat {\mathbf {x}} _ {i + S - j}, t _ {i + S - j}\right) \right]. \tag {42}
$$
In the following, we show that the multi-step loss (4) maximizes a variational lower bound for the maximum likelihood training objective (40).
Theorem A.3. Consider a data set  $\{\mathbf{x}_{0:m}^{(n)}\}_{n=1}^{N}$  with trajectories sampled from SDE (11). Then the multi-step loss (4) maximizes a variational lower bound for the maximum likelihood training objective of the probability flow ODE (40) as  $\Delta t \to 0$ .
Let  $\mu_{t_i}^{\mathrm{ODE}}(\mathbf{x}_{t_j})$  denote the solution of the probability flow ODE (7) integrated backward from time  $t_j$  to  $t_i$  with initial value  $\mathbf{x}_{t_j}$ .
For the maximum likelihood objective, we can derive a variational lower bound
$$
\begin{array}{l} \left. \mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \left[ \log p _ {t _ {i}} ^ {t _ {j}, \mathrm {O D E}} \left(\mathbf {x} _ {t _ {i}}\right) \right] = \mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \left[ \log \left(\mathbb {E} _ {\mathbf {x} _ {t _ {j}}} \left[ p _ {t _ {i}} ^ {t _ {j}, \mathrm {O D E}} \left(\mathbf {x} _ {t _ {i}} \mid \mathbf {x} _ {t _ {j}}\right) \right]\right) \right] \right. (43) \\ = \mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \left[ \log \left(\mathbb {E} _ {\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {t _ {i}}} \left[ \frac {p _ {t _ {i}} \left(\mathbf {x} _ {t _ {i}}\right)}{p _ {t _ {j}} \left(\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {\mathbf {t} _ {\mathbf {i}}}\right)} p _ {t _ {i}} ^ {t _ {j}, \mathrm {O D E}} \left(\mathbf {x} _ {t _ {i}} \mid \mathbf {x} _ {t _ {j}}\right) \right]\right) \right] (44) \\ \geq \mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \mathbb {E} _ {\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {t _ {i}}} \left[ \log \left(\frac {p _ {t _ {j}} \left(\mathbf {x} _ {t _ {j}}\right)}{p _ {t _ {j}} \left(\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {\mathbf {t} _ {\mathbf {i}}}\right)} p _ {t _ {i}} ^ {t _ {j}, \mathrm {O D E}} \left(\mathbf {x} _ {t _ {i}} \mid \mathbf {x} _ {t _ {j}}\right)\right) \right] (45) \\ = \mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \mathbb {E} _ {\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {t _ {i}}} \left[ \log \left(\frac {p _ {t _ {j}} \left(\mathbf {x} _ {t _ {j}}\right)}{p _ {t _ {j}} \left(\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {\mathbf {t} _ {\mathbf {i}}}\right)}\right) + \log \left(p _ {t _ {i}} ^ {t _ {j}, \mathrm {O D E}} \left(\mathbf {x} _ {t _ {i}} \mid \mathbf {x} _ {t _ {j}}\right)\right) \right], (46) \\ \end{array}
$$
where the inequality is due to Jensen's inequality. Since only  $p_{t_i}^{t_j,\mathrm{ODE}}(\mathbf{x}_{t_i}|\mathbf{x}_{t_j})$  depends on  $\theta$ , this is the same as maximizing
$$
\mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \mathbb {E} _ {\mathbf {x} _ {t _ {j}} \mid \mathbf {x} _ {t _ {i}}} \left[ \log \left(p _ {t _ {i}} ^ {t _ {j}, \mathrm {O D E}} \left(\mathbf {x} _ {t _ {i}} \mid \mathbf {x} _ {t _ {j}}\right)\right) \right]. \tag {47}
$$
The probability flow ODE is likelihood-free, which makes it challenging to optimize. Therefore, we relax the objective by perturbing the ODE distributions by convolving them with a Gaussian kernel  $G_{\epsilon}$  with small  $\epsilon > 0$ , see, e.g., Kersting et al. [Ker+20, Gaussian ODE filtering]. This allows us to model the conditional distribution  $p_{t_i}^{t_j,\mathrm{ODE}}|\mathbf{x}_{t_j}$  as a Gaussian distribution with mean  $\mu_{t_i}^{t_j,\mathrm{ODE}}(\mathbf{x}_{t_j})$  and variance  $\sigma^2 = \epsilon$ . Then maximizing (47) reduces to matching the mean of the distribution, i.e., minimizing
$$
\mathbb {E} _ {\mathbf {x} _ {t _ {i}}} \mathbb {E} _ {\mathbf {x} _ {t _ {j}} | \mathbf {x} _ {t _ {i}}} \left[ \left| \left| \mathbf {x} _ {t _ {i}} - \mu_ {t _ {i}} ^ {\mathrm {O D E}} \left(\mathbf {x} _ {t _ {j}}\right) \right| \right| _ {2} ^ {2} \right] \tag {48}
$$
independent of  $\epsilon > 0$ . Since this is true for any time step  $t_j > t_i$  and corresponding simulation state  $\mathbf{x}_{t_j}$  given  $\mathbf{x}_{t_i}$ , we can pick the pairs  $(\mathbf{x}_{t_i}, \mathbf{x}_{t_{i+1}}), (\mathbf{x}_{t_i}, \mathbf{x}_{t_{i+2}}), (\mathbf{x}_{t_i}, \mathbf{x}_{t_{i+3}})$  and so on. Then, we can optimize them jointly by considering the sum of the individual objectives up to a maximum sliding window size
$$
\mathbb {E} _ {\mathbf {x} _ {i: j}} \left[ \sum_ {k = i} ^ {j - 1} \left| \left| \mathbf {x} _ {t _ {k}} - \mu_ {t _ {k}} ^ {\mathrm {O D E}} \left(\mathbf {x} _ {t _ {j}}\right) \right| \right| _ {2} ^ {2} \right]. \tag {49}
$$
As  $\Delta t\to 0$  , we compute the terms  $\mu_{t_k}^{\mathrm{ODE}}(\mathbf{x}_{t_j})$  on a single trajectory starting at  $\mathbf{x}_{t_i}$  with sliding window  $S$  covering the trajectory until  $\mathbf{x}_{t_j}$  via the Euler method, i.e., we can define recursively
$$
\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[ - \mathcal {P} \left(\mathbf {x} _ {i + S - j}\right) + s _ {\theta} \left(\mathbf {x} _ {i + S - j}, t _ {i + S - j}\right) \right]. \tag {50}
$$
By varying the starting points of the sliding window  $\mathbf{x}_{t_i}$ , this yields our proposed multi-step loss  $\mathcal{L}_{\mathrm{multi}}(\theta)$ .
A.3 Denoising Score Matching for Deterministic Simulations
So far, we have considered physical systems that can be modeled by an SDE, i.e., equation (11). While this problem setup is suitable for many scenarios, we would also like to apply a similar
(a) SDE for training data and reverse-time SDE for inference
Figure 7: Variants of training and inference for different physical systems. (a) shows the SDE and reverse-time SDE setup with the Euler-Maruyama discretization when the system is modeled by an SDE. The diffusion term  $g(t)$  is absorbed in the Gaussian random variable  $z_{t} \sim \mathcal{N}(0, g(t)^{2}I)$  and network  $s_{\theta}(\mathbf{x}, t)$ . In (b), we assume that the temporal evolution of the training data is deterministic, i.e., we model the physical system without the diffusion term. However, for inference, we consider the reverse-time SDE of the same form as in (a), where the diffusion coefficient  $g(t)$  is chosen as a hyperparameter that depends on the noise scale added to the data. Then, in (c), we split the Euler step for the backward direction into a physics-only update, adding the Gaussian noise  $z$  and a denoising step by  $s_{\theta}(\mathbf{x}, t)$ .
(b) ODE for training data and SDE for inference
(0E for training data and (c) physics step  $^+$  denoising step for inference
methodology when the system is deterministic, i.e., when we can write the problem as an ordinary stochastic equation
$$
d \mathbf {x} = \mathcal {P} (\mathbf {x}) d t. \tag {51}
$$
In the case of chaotic dynamical systems, this still represents a hard inverse problem, especially when information is lost due to noise added to the trajectories after their generation.
The training setup based on modeling the physics system with an SDE is shown in figure 7a. Figure 7b and 7c illustrate two additional data setup and inference variants for deterministic physical systems modeled by the ODE (51). While for the experiments in sections 3.1 and 3.2 in the main paper, our setup resembles (a), for the buoyancy-driven flow in section 3.3 and the forced isotropic turbulence in section 3.4 in the main paper, we consider (c) as the system is deterministic.
For this variant, the update by  $-\mathcal{P}(\mathbf{x})$  and  $s_{\theta}(\mathbf{x},t)$  is separated into two steps. The temporal evolution from  $t_{i + 1}$  to  $t_i$  is then defined entirely by physics. We apply an additive noise to the system and the update step by  $s_\theta (\mathbf{x},t)$ , which can be interpreted as denoising for a now slightly perturbed state  $\tilde{\mathbf{x}}_{t_i}$ . In this case, we show that the network  $s_{\theta}(\mathbf{x},t)$  still learns the correct score  $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$  during training using denoising score matching. We compare the performance of variants (b) and (c) for the buoyancy-drive flow in appendix E.
When separating physics and score updates, we calculate the updates as
$$
\hat {\mathbf {x}} _ {t _ {i}} = \mathbf {x} _ {t _ {i + 1}} - \Delta t \mathcal {P} \left(\mathbf {x} _ {t _ {i + 1}}\right) \tag {52}
$$
$$
\hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}} = \hat {\mathbf {x}} _ {t _ {i}} + \sqrt {\Delta t} g (t _ {i}) z _ {t _ {i}} \tag {53}
$$
$$
\mathbf {x} _ {t _ {i}} = \hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}} + \Delta t g ^ {2} \left(t _ {i}\right) s _ {\theta} \left(\hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}}, t _ {i}\right), \tag {54}
$$
where  $z_{t_i} \sim \mathcal{N}(0, I)$ . If the physics system is deterministic and  $\Delta t$  is small enough, then we can approximate  $\mathbf{x}_{t_i} \approx \hat{\mathbf{x}}_{t_i}$  and for the moment, we assume that we can write
$$
\hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}} = \mathbf {x} _ {t _ {i}} + \sqrt {\Delta t} g \left(t _ {i}\right) z _ {t _ {i}}. \tag {55}
$$
In this case, we can rewrite the 1-step loss  $\mathcal{L}_{\mathrm{single}}(\theta)$  from (14) to obtain the denoising score matching loss
$$
\mathcal {L} _ {\mathrm {D S M}} (\theta) := \mathbb {E} _ {\left(\mathbf {x} _ {t _ {i}}, \mathbf {x} _ {t _ {i + 1}}\right)} \left[ \left\| \mathbf {x} _ {t _ {i}} - \hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}} - \Delta t g ^ {2} \left(t _ {i}\right) s _ {\theta} \left(\hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}}, t _ {i}\right) \right\| _ {2} ^ {2} \right], \tag {56}
$$
which is the same as minimizing
$$
\mathbb {E} _ {\left(\mathbf {x} _ {t _ {i}}, \mathbf {x} _ {t _ {i + 1}}\right)} \left[ \left| \left| s _ {\theta} \left(\hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}}, t _ {i}\right) - \frac {1}{\Delta t g ^ {2} \left(t _ {i}\right)} \left(\mathbf {x} _ {t _ {i}} - \hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}}\right) \right| \right| _ {2} ^ {2} \right]. \tag {57}
$$
Now, the idea presented in Vincent [Vin11] is that for score matching, we can consider a joint distribution  $p_{t_i}(\mathbf{x}_{t_i},\tilde{\mathbf{x}}_{t_i})$  of sample  $\mathbf{x}_{t_i}$  and corrupted sample  $\tilde{\mathbf{x}}_{t_i}$ . Using Bayes' rule, we can write  $p_{t_i}(\mathbf{x}_{t_i},\tilde{\mathbf{x}}_{t_i}) = p_\sigma (\tilde{\mathbf{x}}_{t_i}|\mathbf{x}_{t_i})p_{t_i}(\mathbf{x}_{t_i})$ . The conditional distribution  $p_{\sigma}(\cdot |\mathbf{x}_{t_i})$  for the corrupted sample is then modeled by a Gaussian with standard deviation  $\sigma = \sqrt{\Delta t} g(t_i)$ , i.e., we can write  $\tilde{\mathbf{x}} = \mathbf{x} + \sqrt{\Delta t} g(t_i)z$  for  $z\sim \mathcal{N}(0,I)$  similar to equation (55). Moreover, we can define the distribution of corrupted data  $q_{\sigma}$  as
$$
q _ {\sigma} (\tilde {\mathbf {x}}) = \int p _ {\sigma} (\tilde {\mathbf {x}} | \mathbf {x}) p _ {t _ {i}} (\mathbf {x}) d \mathbf {x}. \tag {58}
$$
If  $\sigma$  is small, then  $q_{\sigma} \approx p_{t_i}$  and  $\mathrm{KL}(q_{\sigma}||p_{t_i}) \to 0$  as  $\sigma \to 0$ . Importantly, in this case, we can directly compute the score for  $p_{\sigma}(\cdot|\mathbf{x})$  as
$$
\nabla_ {\tilde {\mathbf {x}}} \log p _ {\sigma} (\tilde {\mathbf {x}} | \mathbf {x}) = \frac {1}{\sigma^ {2}} (\mathbf {x} - \tilde {\mathbf {x}}). \tag {59}
$$
Moreover, the theorem proven by Vincent [Vin11] means that we can use the score of the conditional distribution  $p_{\sigma}(\cdot |\mathbf{x})$  to train  $s_{\theta}(\mathbf{x},t)$  to learn the score of  $q_{\sigma}(\mathbf{x})$ , i.e.
$$
\begin{array}{l} \underset {\theta} {\arg \min } \mathbb {E} _ {\tilde {\mathbf {x}} \sim q _ {\theta}} \left[ \left\| s _ {\theta} (\mathbf {x}, t _ {i}) - \nabla_ {\tilde {\mathbf {x}}} \log q _ {\sigma} (\tilde {\mathbf {x}}) \right\| _ {2} ^ {2} \right] (60) \\ = \underset {\theta} {\arg \min } \mathbb {E} _ {\mathbf {x} \sim p _ {t _ {i}}, \tilde {\mathbf {x}} \sim p _ {\sigma (\cdot | \mathbf {x})}} \left[ \| s _ {\theta} (\mathbf {x}, t _ {i}) - \nabla_ {\tilde {\mathbf {x}}} \log p _ {\sigma} (\tilde {\mathbf {x}} | \mathbf {x}) \| _ {2} ^ {2} \right]. (61) \\ \end{array}
$$
By combining (61) and (59), this exactly equals the denoising loss  $\mathcal{L}_{\mathrm{DSM}}(\theta)$  in (57). As  $q_{\sigma} \approx p_{t_i}$ , we also obtain that  $\nabla_{\mathbf{x}} \log q_{\sigma}(\mathbf{x}) \approx \nabla_{\mathbf{x}} \log p_{t_i}(\mathbf{x})$ , so the network  $s_{\theta}(\mathbf{x}, t_i)$  approximately learns the correct score for  $p_{t_i}$ .
We have assumed (55) that the only corruption for  $\hat{\mathbf{x}}_{t_i}^{\mathrm{noise}}$  is the Gaussian noise. This is not true, as we have
$$
\hat {\mathbf {x}} _ {t _ {i}} ^ {\text {n o i s e}} = \mathbf {x} _ {t _ {i}} + \sqrt {\Delta t} g \left(t _ {i}\right) z _ {t _ {i}} + \left(\mathbf {x} _ {t _ {i + 1}} - \Delta t \mathcal {P} \left(\mathbf {x} _ {t _ {i + 1}}\right) - \mathbf {x} _ {t _ {i}}\right), \tag {62}
$$