In this section, we will establish a connection between the mathematical objects that are utilized to construct a narrowly continuous probability path \((\rho_t)_{t\in [0,1]}\) between two distributions \(\rho_0\) and \(\rho_1\). The objects under consideration are the velocity field and the score function, appearing in FM and DM formulations respectively. The following theorem is of main interest.
Theorem 1 . Let \((\Omega, \mathcal F, \mathbb P)\) be a probability space, and \(X_0, X_1\) be two \(\mathbb R^{d}\)-valued independent random variables defined on that probability space such that \(X_0 \sim \rho_0 = \mathcal N(0, I_{d})\) and \(X_1 \sim \rho_1 = p_{\textrm{data}}\). Also, let \((a, b)\) denote an arbitrary FM schedule, and \(v: [0,1] \times \mathbb R^{d} \rightarrow \mathbb R^{d}\) the velocity field that defines the “standard” FM-ODE. For any \(t\in [0,1]\), let \(\rho_t\) denote the density of the stochastic interpolant, \(X_t := a_t \, X_1 + b_t \, X_0\). One has,
\[ v(t, x) = \dfrac{\dot{a}_t}{a_t} \, x + \left( \dfrac{\dot{a}_t}{a_t} - \dfrac{\dot{b}_t}{b_t}\right) \, b_t^{2} \, \nabla_x \log {} \rho_t(x), \ \ (t,x) \in (0,1)\times \mathbb R^{d} \]
To give a formal proof of the preceding theorem, we will make use of the following lemma.
Lemma 1 . Under the assumptions of Theorem 1 denote
\[ D_{\sigma}^{*}(x) = \mathbb E\left[X_1 \mid X_1 + \sigma \, X_0 = x\right], \ \ x\in \mathbb R^{d} \]
the so-called MMSE denoiser of \(X_1\) at level \(\sigma > 0\). One has,
\[ D_{\sigma}^{*}(x) = x + \sigma^{2} \, \nabla_x \log {} \, p_{X_1 + \sigma \, X_0}(x) \]
Proof. One writes,
\[ \begin{align*} \begin{split} p_{X_1 + \sigma \, X_0} (x) & = \int_{\mathbb R^{d}} \, p_{X_1, \, X_1 + \sigma \, X_0} (y, x) \, dy = \int_{\mathbb R^{d}} \, p_{X_1}(y) \, p_{X_1 + \sigma \, X_0 \mid X_1}(x | y) \, dy \\ & = \int_{\mathbb R^{d}} \, p_{X_1}(y) \, (2\pi \, \sigma^{2})^{-d/2} \, \exp{} \left(-||x-y||^{2}/2\sigma^{2}\right) \, dy \\ & := (2\pi \, \sigma^{2})^{-d/2} \, I(x) \end{split} \end{align*} \]
since \(X_1 + \sigma \, X_0 \mid X_1 = y \, \sim \mathcal{N}(y, \sigma^{2} \, I_d)\). Using this equation, we deduce that,
\[ \begin{align*} \begin{split} \nabla_x \log{} p_{X_1 + \sigma \, X_0}(x) & = \nabla_x \log{} I(x) = \dfrac{1}{I(x)} \, \nabla_x I(x) \\ & = \dfrac{1}{I(x)} \, \int_{\mathbb R^{d}} \, p_{X_1}(y) \, \dfrac{(y-x)}{\sigma^{2}} \, \exp{} \left(-||x-y||^{2}/2\sigma^{2}\right) \, dy \\ & = \dfrac{1}{\sigma^{2}} \, \dfrac{1}{I(x)} \, \left[\int_{\mathbb R^{d}} \, y \, p_{X_1}(y) \, \exp{} \left(-||x-y||^{2}/2\sigma^{2}\right) \, dy \, -x \, I(x)\right] \\ & := \dfrac{1}{\sigma^{2}} \, \left(\dfrac{G(x)}{I(x)} - x\right) \end{split} \end{align*} \]
Now observe that,
\[ \begin{align*} \begin{split} D_{\sigma}^{*}(x) & = \dfrac{\int_{\mathbb R^{d}} \, y \, p_{X_1, \, X_1 + \sigma \, X_0} (y, x) \, dy }{\int_{\mathbb R^{d}} \, p_{X_1, \, X_1 + \sigma \, X_0} (y, x) \, dy} = \dfrac{\int_{\mathbb R^{d}} \, y \, p_{X_1}(y) \, p_{X_1 + \sigma \, X_0 \mid X_1}(x|y) \, dy }{\int_{\mathbb R^{d}} p_{X_1}(y) \, p_{X_1 + \sigma \, X_0 \mid X_1}(x|y) \, dy} \\ & = \dfrac{\cancel{(2\pi \, \sigma^{2})^{-d/2}} \, G(x)}{\cancel{(2\pi \, \sigma^{2})^{-d/2}} \, I(x)} \end{split} \end{align*} \]
Therefore,
\[ \nabla_x \log {} p_{X_1 + \sigma \, X_0}(x) = \dfrac{1}{\sigma^{2}} \, \left(D_{\sigma}^{*}(x) - x\right) \]
and a proper rearrangement of the terms in this equation implies the desired result. \(\blacksquare\)
Proof. Theorem 1
Starting by the definition of the velocity field as a conditional expectation we derive the following relation between itself and the MMSE denoiser.
\[ \begin{align*} v(t, x) & = \mathbb E\left[\dot{a}_t \, X_1 + \dot{b}_t \, X_0 \mid a_t \, X_1 + b_t \, X_0 = x\right]\\ & = \mathbb E\left[\dot{a}_t \, X_1 + \dot{b}_t \, \left( \dfrac{x - a_t \, X_1}{b_t} \right) \Bigm \vert a_t \, X_1 + b_t \, X_0 = x\right] \\ & = \dfrac{\dot{b}_t}{b_t} \, x + \left(\dot{a}_t - \dot{b}_t \, \dfrac{a_t}{b_t}\right)\cdot \, \mathbb E\left[ X_1 \Bigm \vert X_1 + \dfrac{b_t}{a_t} \, X_0 = \dfrac{x}{a_t} \right] \\ & = \dfrac{\dot{b}_t}{b_t} \, x + \left(\dot{a}_t - \dot{b}_t \, \dfrac{a_t}{b_t}\right)\cdot D_{b_t/a_t}^{*}(x/a_t) \end{align*} \tag{1}\]
By Lemma 1, we have,
\[ D_{b_t/a_t}^{*}(x/a_t) = \dfrac{x}{a_t} + \dfrac{b_t^{2}}{a_t^{2}} \, \nabla_{x/a_t} \log{} p_{X_1 + \frac{b_t}{a_t} \, X_0} (x/a_t) \tag{2}\]
Notice that,
\[ p_{X_1 + \frac{b_t}{a_t} \, X_0}(x/a_t) = p_{\frac{X_t}{a_t}}(x/a_t) = a_t^{d} \, \rho_t(x) \]
Taking the logarithm, differentiating with respect to \(x/a_t\), and substituting in Equation 2 we get,
\[ D_{b_t/a_t}^{*}(x/a_t) = \dfrac{x}{a_t} + \dfrac{b_t^{2}}{a_t^{2}} \, \nabla_{x/a_t} \log{} \rho_t(x) = \dfrac{x}{a_t} + \dfrac{b_t^{2}}{a_t} \, \nabla_x \log {} \rho_t(x) \tag{3}\]
Substituting Equation 3 in Equation 1, one writes,
\[ v(t, x) = \dfrac{\dot{a}_t}{a_t} \, x + \left( \dfrac{\dot{a}_t}{a_t} - \dfrac{\dot{b}_t}{b_t} \right) \, b_t^{2} \, \nabla_x \log{} \rho_t(x), \ \ (t,x) \in (0,1) \times \mathbb R^{d} \]
which is the desired expression. \(\blacksquare\)