Appendix D: Supplementary Materials and Proofs

Variational Perspective

Theorem in Chapter 2 (Equivalence Between Marginal and Conditional KL Minimization) Equivalence Between Marginal and Conditional KL Minimization

Proof.

Derivation of Equation in Chapter 2.

We start by expanding the right-hand side expectation: \[\begin{align*} &\mathbb{E}_{p({\mathbf{x}}_0, {\mathbf{x}}_i)}\left[ \mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \| p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) \right] \\ =& \int \int p({\mathbf{x}}_0, {\mathbf{x}}_i) \mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \| p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \mathop{}\!\mathrm{d}{\mathbf{x}}_i. \end{align*}\] By the definition of KL divergence, \[\begin{align*} &\mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \| p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) = \int p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \mathop{}\!\mathrm{d}{\mathbf{x}}_{i-1}. \end{align*}\] Substituting this into the expectation, we have \[\begin{align*} &\int \int \int p({\mathbf{x}}_0, {\mathbf{x}}_i) p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \mathop{}\!\mathrm{d}{\mathbf{x}}_{i-1} \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \mathop{}\!\mathrm{d}{\mathbf{x}}_i. \end{align*}\] Using the chain rule of probability, \[\begin{align*} p({\mathbf{x}}_0, {\mathbf{x}}_i) &= p({\mathbf{x}}_i) p({\mathbf{x}}_0|{\mathbf{x}}_i), \end{align*}\] we rewrite the integral as \[\begin{align*} &\int p({\mathbf{x}}_i) \int p({\mathbf{x}}_0|{\mathbf{x}}_i) \int p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \mathop{}\!\mathrm{d}{\mathbf{x}}_{i-1} \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \mathop{}\!\mathrm{d}{\mathbf{x}}_i. \end{align*}\] This allows us to express the expectation in nested form: \[\begin{align*} &\mathbb{E}_{p({\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)} \left[ \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \right] \right] \right]. \end{align*}\]

Next, we apply the decomposition of the logarithm: \[\begin{align*} \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} =\log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} + \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)}. \end{align*}\] Substituting this back into the expectation gives two terms: \[\begin{align*} &\mathbb{E}_{p({\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)} \left[ \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)}{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \right] \right] \right] \\ &+ \mathbb{E}_{p({\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)} \left[ \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \right] \right] \right]. \end{align*}\] Since the second logarithmic term does not depend on \({\mathbf{x}}_0\), by the law of total probability \[\begin{align*} \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0)} \left[ \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \right] \right] = \mathbb{E}_{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \left[ \log \frac{p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)}{p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)} \right]. \end{align*}\] Similarly, the first term is the KL divergence \[\begin{align*} \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ \mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \| p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) \right]. \end{align*}\] Putting it all together, we obtain the decomposition: \[\begin{align*} &\mathbb{E}_{p({\mathbf{x}}_0, {\mathbf{x}}_i)} \left[ \mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \| p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) \right] \\ =& \mathbb{E}_{p({\mathbf{x}}_i)} \left[ \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ \mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i, {\mathbf{x}}_0) \| p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) \right] \right] \\ &+ \mathbb{E}_{p({\mathbf{x}}_i)} \left[ \mathcal{D}_{\text{KL}} \big(p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i) \| p_\phi({\mathbf{x}}_{i-1}|{\mathbf{x}}_i)\big) \right]. \end{align*}\]

Proof of Optimality. To prove: \[p^*({\mathbf{x}}_{i-1}|{\mathbf{x}}_i) = p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i) = \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_i)} \left[ p({\mathbf{x}}_{i-1}|{\mathbf{x}}_i,{\mathbf{x}}_0) \right], \qquad {\mathbf{x}}_i \sim p_i.\] The first identity follows from the fact that the KL divergence \(\mathcal{D}_{\mathrm{KL}}(p \| p_{\bm{\phi}})\) is minimized when \(p^* = p\), assuming the parameterization is sufficiently expressive. The second identity follows directly from the law of total probability. \(\blacksquare\)

Theorem in Chapter 2 (DDPM’s ELBO) ELBO of Diffusion Model

Proof.

For notational simplicity, we denote \[{\mathbf{x}}_{1:L}:=({\mathbf{x}}_1,\ldots,{\mathbf{x}}_L).\]

Step 1: Apply Jensen’s Inequality. The marginal log-likelihood is \[\log p_{\bm{\phi}}({\mathbf{x}}_0) = \log \int p_{\bm{\phi}}({\mathbf{x}}_0,{\mathbf{x}}_{1:L})\,\mathop{}\!\mathrm{d}{\mathbf{x}}_{1:L},\] where the joint generative distribution is \[p_{\bm{\phi}}({\mathbf{x}}_0,{\mathbf{x}}_{1:L}) = p_{\text{prior}}({\mathbf{x}}_L)\prod_{i=1}^L p_{\bm{\phi}}({\mathbf{x}}_{i-1}| {\mathbf{x}}_i).\]

We introduce the forward noising distribution \[p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0) = \prod_{i=1}^L p({\mathbf{x}}_i| {\mathbf{x}}_{i-1}),\] and rewrite \[\log p_{\bm{\phi}}({\mathbf{x}}_0) = \log \int p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0) \frac{p_{\bm{\phi}}({\mathbf{x}}_0,{\mathbf{x}}_{1:L})}{p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0)} \,\mathop{}\!\mathrm{d}{\mathbf{x}}_{1:L}.\] Applying Jensen’s inequality yields \[\log p_{\bm{\phi}}({\mathbf{x}}_0) \ge \mathbb{E}_{p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0)} \left[ \log \frac{p_{\bm{\phi}}({\mathbf{x}}_0,{\mathbf{x}}_{1:L})}{p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0)} \right] =: \mathcal{L}_{\text{ELBO}}({\mathbf{x}}_0;\bm{\phi}),\] and hence \[-\log p_{\bm{\phi}}({\mathbf{x}}_0) \le -\mathcal{L}_{\text{ELBO}}({\mathbf{x}}_0;\bm{\phi}).\]

Step 2: Expand the ELBO. Substituting the factorizations into the ELBO gives \[\begin{align*} \mathcal{L}_{\text{ELBO}}({\mathbf{x}}_0;\bm{\phi}) &= \mathbb{E}_{p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0)} \Bigg[ \log p_{\text{prior}}({\mathbf{x}}_L) + \sum_{i=1}^L \log p_{\bm{\phi}}({\mathbf{x}}_{i-1}| {\mathbf{x}}_i) - \sum_{i=1}^L \log p({\mathbf{x}}_i| {\mathbf{x}}_{i-1}) \Bigg]. \end{align*}\] Therefore, \[\begin{align*} -\mathcal{L}_{\text{ELBO}}({\mathbf{x}}_0;\bm{\phi}) &= \mathbb{E}_{p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0)} \Bigg[ \log \frac{p({\mathbf{x}}_L| {\mathbf{x}}_0)}{p_{\text{prior}}({\mathbf{x}}_L)} + \sum_{i=2}^L \log \frac{p({\mathbf{x}}_{i-1}| {\mathbf{x}}_i,{\mathbf{x}}_0)}{p_{\bm{\phi}}({\mathbf{x}}_{i-1}| {\mathbf{x}}_i)} - \log p_{\bm{\phi}}({\mathbf{x}}_0| {\mathbf{x}}_1) \Bigg]. \end{align*}\] Here we used the identity \[p({\mathbf{x}}_{1:L}| {\mathbf{x}}_0) = p({\mathbf{x}}_L| {\mathbf{x}}_0)\prod_{i=2}^L p({\mathbf{x}}_{i-1}| {\mathbf{x}}_i,{\mathbf{x}}_0),\] which follows from the Markov structure of the forward process.

Now group the terms according to their dependence: \[\begin{align*} -\mathcal{L}_{\text{ELBO}}({\mathbf{x}}_0;\bm{\phi}) &= \underbrace{ \mathbb{E}_{p({\mathbf{x}}_L| {\mathbf{x}}_0)} \left[ \log \frac{p({\mathbf{x}}_L| {\mathbf{x}}_0)}{p_{\text{prior}}({\mathbf{x}}_L)} \right] }_{\mathcal{L}_{\text{prior}}({\mathbf{x}}_0)} + \underbrace{ \mathbb{E}_{p({\mathbf{x}}_1| {\mathbf{x}}_0)} \left[ -\log p_{\bm{\phi}}({\mathbf{x}}_0| {\mathbf{x}}_1) \right] }_{\mathcal{L}_{\text{recon.}}({\mathbf{x}}_0;\bm{\phi})} \\ &\quad+ \underbrace{ \sum_{i=2}^L \mathbb{E}_{p({\mathbf{x}}_i| {\mathbf{x}}_0)} \left[ \mathcal{D}_{\mathrm{KL}} \big( p({\mathbf{x}}_{i-1}| {\mathbf{x}}_i,{\mathbf{x}}_0) \| p_{\bm{\phi}}({\mathbf{x}}_{i-1}| {\mathbf{x}}_i) \big) \right] }_{\mathcal{L}_{\text{diffusion}}({\mathbf{x}}_0;\bm{\phi})}. \end{align*}\] This proves the claimed decomposition. \(\blacksquare\)

Score-Based Perspective

Proposition in Chapter 3 (Hyvärinen’s Tractable Form of SM): Tractable Score Matching via Integration by Parts

Proof.

Expanding \(\mathcal{L}_{\text{SM}}(\bm{\phi})\). Let us expand the squared difference inside the expectation: \[\begin{align*} \mathcal{L}_{\text{SM}}(\bm{\phi}) &= \frac{1}{2}\mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \left\lVert{\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})\right\rVert_2^2 - 2 \langle {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}}), {\mathbf{s}}({\mathbf{x}}) \rangle + \left\lVert{\mathbf{s}}({\mathbf{x}})\right\rVert_2^2 \right] \\ &= \frac{1}{2}\mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \left\lVert{\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})\right\rVert_2^2 \right] - \mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \langle {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}}), {\mathbf{s}}({\mathbf{x}}) \rangle \right] \\ &\quad + \frac{1}{2}\mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \left\lVert{\mathbf{s}}({\mathbf{x}})\right\rVert_2^2 \right]. \end{align*}\] We now focus on the cross-product term: \[\mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \langle {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}}), {\mathbf{s}}({\mathbf{x}}) \rangle \right].\] Using the fact that \[\nabla_{{\mathbf{x}}} \log p_{\text{data}}({\mathbf{x}}) = \frac{\nabla_{{\mathbf{x}}}p_{\text{data}}({\mathbf{x}})}{p_{\text{data}}({\mathbf{x}})},\] and assuming \(p_{\text{data}}({\mathbf{x}})\) is not zero (e.g., on its support), the cross-product term becomes: \[\begin{align*} \mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \langle {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}}), {\mathbf{s}}({\mathbf{x}}) \rangle \right] &= \int {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})^\top \nabla_{{\mathbf{x}}} \log p_{\text{data}}({\mathbf{x}}) p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}\\ &= \int {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})^\top \nabla_{{\mathbf{x}}}p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}\\ &= \sum_{i=1}^{D} \int s_{\bm{\phi}}^{(i)}({\mathbf{x}}) \partial_{x_i}p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}, \end{align*}\] where \(s_{\bm{\phi}}^{(i)}({\mathbf{x}})\) is the \(i\)-th component of the score function \[{\mathbf{s}}_{\bm{\phi}} = \left( s_{\bm{\phi}}^{(1)}, s_{\bm{\phi}}^{(2)}, \dots, s_{\bm{\phi}}^{(D)} \right).\] We use the following integration-by-parts formula (Evans 2010), which is derived from standard calculus:

We apply this formula to \(u({\mathbf{x}}) := s_{\bm{\phi}}^{(i)}({\mathbf{x}})\) and \(v({\mathbf{x}}) = p_{\text{data}}({\mathbf{x}})\) for all \(i = 1, \dots, D\), assuming that \[|u({\mathbf{x}}) v({\mathbf{x}})| \to 0 \quad \text{as } R \to \infty.\] Summing the results over all \(i = 1, \dots, D\), we get: \[\begin{align*} \mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \langle {\mathbf{s}}_{\bm{\phi}}({\mathbf{x}}), {\mathbf{s}}({\mathbf{x}}) \rangle \right] &= - \sum_{i=1}^{D} \int \partial_{x_i}s_{\bm{\phi}}^{(i)}({\mathbf{x}}) p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}\\ &= - \mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[\mathop{\mathrm{Tr}}\left(\nabla_{{\mathbf{x}}}{\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})\right) \right]. \end{align*}\] Combining all results, we have: \[\begin{align*} \mathcal{L}_{\text{SM}}(\bm{\phi}) &= \underbrace{\mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[\mathop{\mathrm{Tr}}\left(\nabla_{{\mathbf{x}}}{\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})\right) + \frac{1}{2} \left\lVert{\mathbf{s}}_{\bm{\phi}}({\mathbf{x}})\right\rVert_2^2 \right]}_{\widetilde{\mathcal{L}}_{\text{SM}}(\bm{\phi})} \\ &\quad + \underbrace{\frac{1}{2}\mathbb{E}_{{\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})} \left[ \left\lVert{\mathbf{s}}({\mathbf{x}})\right\rVert_2^2 \right]}_{=:C}, \end{align*}\] where \(C\) depends only on the distribution \(p_{\text{data}}\), which concludes the proof. \(\blacksquare\)

Theorem in Chapter 3 (Equivalence of L_SM and L_DSM) Equivalence Between SM and DSM Minimization

Proof.

Expanding both \(\mathcal{L}_{\text{SM}}(\bm{\phi}; \sigma)\) and \(\mathcal{L}_{\text{DSM}}(\bm{\phi}; \sigma)\), we have: \[\begin{align*} &\mathcal{L}_{\text{SM}}(\bm{\phi}; \sigma) = \frac{1}{2} \mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})} \begin{aligned}[t] \Big[ \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2 &- 2 {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}}) \\ &+ \left\lVert\nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}})\right\rVert_2^2 \Big], \end{aligned} \\ &\mathcal{L}_{\text{DSM}}(\bm{\phi}; \sigma) = \frac{1}{2} \mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})} \begin{aligned}[t] \Big[ \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2 &- 2 {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \log p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}}) \\ &+ \left\lVert\nabla_{\tilde{{\mathbf{x}}}} \log p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})\right\rVert_2^2 \Big]. \end{aligned} \end{align*}\] Subtracting the two equations yields: \[\begin{align*} &\mathcal{L}_{\text{SM}}(\bm{\phi}; \sigma) -\mathcal{L}_{\text{DSM}}(\bm{\phi}; \sigma) \\= \frac{1}{2} &\Bigg( \mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})} \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2 - \mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})} \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2\Bigg) \\ \quad - & \begin{aligned}[t] \Bigg( \mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})} &\left[ {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}}) \right] \\ &- \mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})}\left[ {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}}| {\mathbf{x}})\right] \Bigg) \end{aligned} \\ \quad + \frac{1}{2}&\Bigg(\mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})} \left\lVert\nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}})\right\rVert_2^2 - \mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})} \left\lVert\nabla_{\tilde{{\mathbf{x}}}} \log p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})\right\rVert_2^2 \Bigg). \end{align*}\] Next, we address the three terms one at a time. For the first term, since \(p_\sigma(\tilde{{\mathbf{x}}}) = \int p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}}) p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}\), we can rewrite it as: \[\begin{align*} \mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})} \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2 &= \int \Big( \int p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}}) p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}\Big) \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2 \mathop{}\!\mathrm{d}\tilde{{\mathbf{x}}} \\ &= \int p_{\text{data}}({\mathbf{x}}) \int p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}}) \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2 \mathop{}\!\mathrm{d}\tilde{{\mathbf{x}}} \mathop{}\!\mathrm{d}{\mathbf{x}} \\ &= \mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})} \left\lVert{\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)\right\rVert_2^2. \end{align*}\] Thus, the first term is zero. For the second term: \[\begin{align} \label{eq:dsm-sec-term} \begin{aligned} &\mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})}\left[ {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}}) \right] \\=&\int p_\sigma(\tilde{{\mathbf{x}}}) {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \frac{\nabla_{\tilde{{\mathbf{x}}}} p_\sigma(\tilde{{\mathbf{x}}})}{p_\sigma(\tilde{{\mathbf{x}}})} \mathop{}\!\mathrm{d}\tilde{\mathbf{x}} \\=&\int {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \int p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}}) p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}{\mathbf{x}}\mathop{}\!\mathrm{d}\tilde{\mathbf{x}} \\=&\int \int {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}}p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}}) p_{\text{data}}({\mathbf{x}}) \mathop{}\!\mathrm{d}\tilde{\mathbf{x}}\mathop{}\!\mathrm{d}{\mathbf{x}} \\= &\mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})}\left[ {\mathbf{s}}_{\bm{\phi}}(\tilde{{\mathbf{x}}}; \sigma)^\top \nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}}| {\mathbf{x}})\right]. \end{aligned} \end{align}\] Thus, it is also zero. For the third term, note that: \[C:=\frac{1}{2}\Bigg(\mathbb{E}_{\tilde{{\mathbf{x}}} \sim p_\sigma(\tilde{{\mathbf{x}}})} \left\lVert\nabla_{\tilde{{\mathbf{x}}}} \log p_\sigma(\tilde{{\mathbf{x}}})\right\rVert_2^2 - \mathbb{E}_{ p_{\text{data}}({\mathbf{x}})p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})} \left\lVert\nabla_{\tilde{{\mathbf{x}}}} \log p_{\sigma}(\tilde{{\mathbf{x}}} | {\mathbf{x}})\right\rVert_2^2 \Bigg)\] depends only on \(p_{\text{data}}({\mathbf{x}})\) and \(p_{\sigma}(\tilde{{\mathbf{x}}}|{\mathbf{x}})\), and hence it is constant with respect to \(\bm{\phi}\). \(\blacksquare\)

Lemma in Chapter 3 (Tweedie’s Formula): Tweedie’s Formula

We first state a more general form of Tweedie’s formula, which considers time-dependent Gaussian perturbations, and we provide its proof below.

Tweedie’s Identity with Time-Dependent Parameters.

Let \({\mathbf{x}}_t|{\mathbf{x}}_0 \sim \mathcal{N}\big(\cdot; \alpha_t{\mathbf{x}}_0, \sigma_t^2{\mathbf{I}}\big)\). Then Tweedie’s formula says \[\begin{align*} \alpha_t \mathbb{E}_{\mathbf{x}_0 \sim p(\mathbf{x}_0|{\mathbf{x}}_t)}[\mathbf{x}_0|{\mathbf{x}}_t] = {\mathbf{x}}_t + \sigma_t^2 \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t), \end{align*}\] where the expectation is taken over the posterior distribution \(p(\mathbf{x}_0|{\mathbf{x}}_t)\) of \(\mathbf{x}_0\) given the observed \({\mathbf{x}}_t\), and \(p_t({\mathbf{x}}_t)\) is the marginal density of \({\mathbf{x}}_t\).

Proof.

Marginal Density and Its Score.

We recall that the marginal density of \({\mathbf{x}}_t\) is given by \[\begin{align*} p_t({\mathbf{x}}_t) = \int p_t({\mathbf{x}}_t|{\mathbf{x}}_0) p_0({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0. \end{align*}\] We now compute the score function: \[\begin{align*} \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t) = \frac{\nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t)}{p_t({\mathbf{x}}_t)} = \frac{1}{p_t({\mathbf{x}}_t)} \int \nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t|{\mathbf{x}}_0) p_0({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0. \end{align*}\] We therefore need to compute the gradient of the conditional density.

Gradient of the Conditional and Rearrangement.

The gradient of the conditional Gaussian density is: \[\begin{align*} \nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t|{\mathbf{x}}_0) = -p_t({\mathbf{x}}_t|{\mathbf{x}}_0) \cdot \sigma_t^{-2} ({\mathbf{x}}_t - \alpha_t {\mathbf{x}}_0). \end{align*}\] Substituting this into the previous expression, we have: \[\begin{align*} \nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t) &= \int \nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t|{\mathbf{x}}_0) p_0({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \\ &= -\sigma_t^{-2} \int ({\mathbf{x}}_t - \alpha_t {\mathbf{x}}_0) p_t({\mathbf{x}}_t|{\mathbf{x}}_0) p_0({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \\ &= -\sigma_t^{-2} \int ({\mathbf{x}}_t - \alpha_t {\mathbf{x}}_0) p({\mathbf{x}}_0|{\mathbf{x}}_t) p_t({\mathbf{x}}_t) \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \\ &= -p_t({\mathbf{x}}_t) \sigma_t^{-2} \left({\mathbf{x}}_t - \alpha_t \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}[{\mathbf{x}}_0\vert {\mathbf{x}}_t] \right). \end{align*}\] Dividing both sides by \(p_t({\mathbf{x}}_t)\), we obtain: \[\begin{align*} \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t) = -\sigma_t^{-2} \left({\mathbf{x}}_t - \alpha_t \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}[{\mathbf{x}}_0\vert {\mathbf{x}}_t] \right). \end{align*}\] Rearranging yields: \[\begin{align*} {\mathbf{x}}_t + \sigma_t^{2} \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t) = \alpha_t \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}[{\mathbf{x}}_0\vert {\mathbf{x}}_t]. \end{align*}\] This completes the derivation.

Stein’s Identity and Surrogate SURE Objective

Stein’s Identity.

Stein’s identity is the integration-by-parts technique that turns expectations under an unknown density into expectations of observable functions and their derivatives, which cancels the partition function and enables unbiased, tractable objectives and tests without ever evaluating the unknown density or the partition function. We begin with the simplest one-dimensional case and then extend it to the form needed to prove the surrogate loss for SURE.

1D, Standard Normal Case.

If \(z \sim \mathcal{N}(0,1)\) and \(f\) has suitable decay, then Stein’s identity states: \[\mathbb{E}[f'(z)] = \mathbb{E}[z f(z)] .\] Denote \(\phi(z) := \tfrac{1}{\sqrt{2\pi}} e^{-z^2/2}\), the one-dimensional standard normal density. The proof follows by integration by parts, using \(\phi'(z) = -z\phi(z)\), together with the vanishing boundary term. To see this precisely, we compute \[\mathbb{E}[f'(z)] = \int_{-\infty}^{\infty} f'(z) \phi(z) \mathop{}\!\mathrm{d}z.\] By integration by parts, with \(u=\phi(z)\) and \(\mathop{}\!\mathrm{d}v=f'(z) \mathop{}\!\mathrm{d}z\), we obtain \[\int f'(z) \phi(z) \mathop{}\!\mathrm{d}z = \Big[ f(z) \phi(z) \Big]_{-\infty}^{\infty} - \int f(z) \phi'(z) \mathop{}\!\mathrm{d}z.\] Since \(\phi'(z) = -z\phi(z)\) and \(f(z)\phi(z)\to 0\) as \(|z|\to\infty\) (decay condition), the boundary term vanishes and we have \[\mathbb{E}[f'(z)] = \int f(z) z \phi(z) \mathop{}\!\mathrm{d}z = \mathbb{E}[z f(z)].\] This completes the derivation and proves the one-dimensional case of Stein’s identity.

Multivariate, Standard Normal Case.

If \({\mathbf{z}}\sim \mathcal{N}(\mathbf{0},{\mathbf{I}}_D)\) and \(g:\mathbb{R}^D \to \mathbb{R}\), then Stein’s identity is \[\mathbb{E}[\nabla g({\mathbf{z}})] = \mathbb{E}[{\mathbf{z}}g({\mathbf{z}})].\] Equivalently, for \(\mathbf{u}:\mathbb{R}^D \to \mathbb{R}^D\), \[\begin{align} \label{eq:stein-multi} \mathbb{E}[ \nabla_{{\mathbf{z}}} \cdot\mathbf{u}({\mathbf{z}})] = \mathbb{E}[{\mathbf{z}}^\top \mathbf{u}({\mathbf{z}})] . \end{align}\]

Identity Needed for SURE.

With \(\tilde{{\mathbf{x}}}={\mathbf{x}}+\sigma{\mathbf{z}}\), where \({\mathbf{z}}\sim \mathcal{N}(\mathbf{0},{\mathbf{I}}_D)\), and any vector function \(\mathbf{a}\) of suitable regularity, \[\begin{align} \label{eq:stein-for-sure} \mathbb{E}\big[(\tilde{{\mathbf{x}}}-{\mathbf{x}})^\top \mathbf{a}(\tilde{{\mathbf{x}}}) \big| {\mathbf{x}}\big] = \sigma^2 \mathbb{E}\big[ \nabla_{\tilde{{\mathbf{x}}}} \cdot\mathbf{a}(\tilde{{\mathbf{x}}}) \big| {\mathbf{x}}\big]. \end{align}\] This is obtained by applying Equation \(\eqref{eq:stein-multi}\) and using the chain rule.

Deriving SURE from the Conditional MSE.

Let \({\mathbf{D}}:\mathbb{R}^D \to \mathbb{R}^D\) be a denoiser and define \[R({\mathbf{D}};{\mathbf{x}}):=\mathbb{E}\left[\|{\mathbf{D}}(\tilde{{\mathbf{x}}})-{\mathbf{x}}\|_2^2| {\mathbf{x}}\right].\] Expand around \(\tilde{{\mathbf{x}}}\): \[\begin{aligned} &~R({\mathbf{D}};{\mathbf{x}}) \\ =&~\mathbb{E}\left[\|{\mathbf{D}}(\tilde{{\mathbf{x}}})-\tilde{{\mathbf{x}}}\|^2 \big| {\mathbf{x}}\right] +2 \mathbb{E}\left[({\mathbf{D}}(\tilde{{\mathbf{x}}})-\tilde{{\mathbf{x}}})^\top(\tilde{{\mathbf{x}}}-{\mathbf{x}}) \big| {\mathbf{x}}\right] +\mathbb{E}\left[\|\tilde{{\mathbf{x}}}-{\mathbf{x}}\|^2 \big| {\mathbf{x}}\right] \\ =&~\mathbb{E}\left[\|{\mathbf{D}}(\tilde{{\mathbf{x}}})-\tilde{{\mathbf{x}}}\|^2 \big| {\mathbf{x}}\right] +2\Big(\underbrace{\mathbb{E}[(\tilde{{\mathbf{x}}}-{\mathbf{x}})^\top {\mathbf{D}}(\tilde{{\mathbf{x}}})|{\mathbf{x}}]}_{\substack{\sigma^2 \mathbb{E}[ \nabla_{\tilde{{\mathbf{x}}}} \cdot{\mathbf{D}}(\tilde{{\mathbf{x}}})|{\mathbf{x}}] \\ \text{by (\ref{eq:stein-for-sure})} }} -\underbrace{\mathbb{E}[(\tilde{{\mathbf{x}}}-{\mathbf{x}})^\top \tilde{{\mathbf{x}}}|{\mathbf{x}}]}_{ \sigma^2 D}\Big) \\&\qquad\qquad\qquad+\underbrace{\mathbb{E}[\|\tilde{{\mathbf{x}}}-{\mathbf{x}}\|^2|{\mathbf{x}}]}_{ \sigma^2 D} \\ =&~\mathbb{E}\left[\|{\mathbf{D}}(\tilde{{\mathbf{x}}})-\tilde{{\mathbf{x}}}\|^2 +2\sigma^2 \nabla_{\tilde{{\mathbf{x}}}} \cdot{\mathbf{D}}(\tilde{{\mathbf{x}}}) -D\sigma^2 | {\mathbf{x}}\right]. \end{aligned}\] Therefore the observable surrogate \[\mathrm{SURE}({\mathbf{D}};\tilde{{\mathbf{x}}}) := \|{\mathbf{D}}(\tilde{{\mathbf{x}}})-\tilde{{\mathbf{x}}}\|_2^2 +2\sigma^2 \nabla_{\tilde{{\mathbf{x}}}} \cdot{\mathbf{D}}(\tilde{{\mathbf{x}}}) -D\sigma^2\] satisfies \(\mathbb{E}\left[\mathrm{SURE}({\mathbf{D}};\tilde{{\mathbf{x}}}) \big| {\mathbf{x}}\right]=R({\mathbf{D}};{\mathbf{x}})\). Minimizing SURE (in expectation or empirically) is thus equivalent to minimizing the true conditional MSE while using only noisy observations.

Theorem in Chapter 4 (Marginal Alignment via the Fokker–Planck Equation) Marginal Alignment via Fokker–Planck

Proof.

Part 1: Fokker-Planck Equation for the Forward SDE.

Consider the forward SDE: \[\mathop{}\!\mathrm{d}{\mathbf{x}}(t) = {\mathbf{f}}({\mathbf{x}}(t), t) \mathop{}\!\mathrm{d}t + g(t) \mathop{}\!\mathrm{d}{\mathbf{w}}(t).\] The diffusion matrix is \(\sigma(t) = g(t) I_D\), so \(\sigma(t) \sigma(t)^T = g^2(t) I_D\). The Fokker-Planck equation for the marginal density \(p_t({\mathbf{x}})\) of \({\mathbf{x}}(t)\) is: \[\partial_t p_t({\mathbf{x}}) = -\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{f}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr] + \frac{1}{2} \sum_{i,j=1}^D \frac{\partial^2}{\partial x_i \partial x_j} \bigl[ (g^2(t) \delta_{ij}) p_t({\mathbf{x}}) \bigr].\] Compute the diffusion term: \[\sum_{i,j=1}^D \frac{\partial^2}{\partial x_i \partial x_j} \bigl[ g^2(t) \delta_{ij} p_t({\mathbf{x}}) \bigr] = \sum_{i=1}^D \frac{\partial^2}{\partial x_i^2} \bigl[ g^2(t) p_t({\mathbf{x}}) \bigr] = g^2(t) \Delta_{{\mathbf{x}}} p_t({\mathbf{x}}).\] Thus: \[\partial_t p_t({\mathbf{x}}) = -\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{f}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr] + \frac{1}{2} g^2(t) \Delta_{{\mathbf{x}}} p_t({\mathbf{x}}).\] Now, rewrite using: \[{\mathbf{v}}({\mathbf{x}}, t) = {\mathbf{f}}({\mathbf{x}}, t) - \frac{1}{2} g^2(t) \nabla_{{\mathbf{x}}} \log p_t({\mathbf{x}}).\] Since \(\nabla_{{\mathbf{x}}} \log p_t({\mathbf{x}}) = \frac{\nabla_{{\mathbf{x}}} p_t({\mathbf{x}})}{p_t({\mathbf{x}})}\), compute: \[\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr] = \nabla_{{\mathbf{x}}} \cdot \biggl[ {\mathbf{f}}({\mathbf{x}}, t) p_t({\mathbf{x}}) - \frac{1}{2} g^2(t) \frac{\nabla_{{\mathbf{x}}} p_t({\mathbf{x}})}{p_t({\mathbf{x}})} p_t({\mathbf{x}}) \biggr].\] The second term is: \[\nabla_{{\mathbf{x}}} \cdot \biggl[ -\frac{1}{2} g^2(t) \nabla_{{\mathbf{x}}} p_t({\mathbf{x}}) \biggr] = -\frac{1}{2} g^2(t) \Delta_{{\mathbf{x}}} p_t({\mathbf{x}}).\] Thus: \[\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr] = \nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{f}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr] - \frac{1}{2} g^2(t) \Delta_{{\mathbf{x}}} p_t({\mathbf{x}}).\] Therefore: \[\partial_t p_t({\mathbf{x}}) = -\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr],\] verifying the Fokker-Planck equation in both forms.

Part 2: PF-ODE Marginal Densities.

Consider the PF-ODE: \[\frac{\mathop{}\!\mathrm{d}\tilde{\mathbf{x}}(t)}{\mathop{}\!\mathrm{d}t} = {\mathbf{v}}(\tilde{\mathbf{x}}(t), t), \quad {\mathbf{v}}({\mathbf{x}}, t) = {\mathbf{f}}({\mathbf{x}}, t) - \frac{1}{2} g^2(t) \nabla_{{\mathbf{x}}} \log p_t({\mathbf{x}}).\]

Forward Direction: \(\tilde{\mathbf{x}}(0) \sim p_0\). Let \(\tilde{\mathbf{x}}(t)\) follow the PF-ODE with \(\tilde{\mathbf{x}}(0) \sim p_0\). The flow map \(\bm{\Psi}_t: \mathbb{R}^D \to \mathbb{R}^D\) is defined by: \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} \bm{\Psi}_t(\mathbf{x}_0) = {\mathbf{v}}(\bm{\Psi}_t(\mathbf{x}_0), t), \quad \bm{\Psi}_0(\mathbf{x}_0) = \mathbf{x}_0.\] Since \(\tilde{\mathbf{x}}(t) = \bm{\Psi}_t(\tilde{\mathbf{x}}(0))\), the density \(\tilde p_t({\mathbf{x}})\) of \(\tilde{\mathbf{x}}(t)\) satisfies the continuity equation: \[\partial_t \tilde p_t({\mathbf{x}}) = -\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, t) \tilde p_t({\mathbf{x}}) \bigr].\] Since \(\tilde{\mathbf{x}}(0) \sim p_0\), we have \(\tilde p_0({\mathbf{x}}) = p_0({\mathbf{x}})\). From Part 1, \(p_t({\mathbf{x}})\) satisfies: \[\partial_t p_t({\mathbf{x}}) = -\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, t) p_t({\mathbf{x}}) \bigr].\] Both \(\tilde p_t\) and \(p_t\) satisfy the same continuity equation with the same initial condition \(p_0\). Assuming sufficient smoothness (e.g., \({\mathbf{v}}\in C^1\)), the solution is unique in some appropriate function space, so \(\tilde p_t = p_t\). Thus, \(\tilde{\mathbf{x}}(t) \sim p_t\) for all \(t \in [0, T]\).

Backward Direction: \(\tilde{\mathbf{x}}(T) \sim p_T\). Now, let \(\tilde{\mathbf{x}}(t)\) follow the PF-ODE backward from \(t = T\) to \(t = 0\), with \(\tilde{\mathbf{x}}(T) \sim p_T\). The ODE is: \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} \tilde{\mathbf{x}}(t) = {\mathbf{v}}(\tilde{\mathbf{x}}(t), t).\] Let \(s = T - t\), so the backward ODE becomes: \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s} \tilde{\mathbf{x}}(T - s) = -{\mathbf{v}}(\tilde{\mathbf{x}}(T - s), T - s).\] The density \(\tilde p_{T-s}({\mathbf{x}})\) of \(\tilde{\mathbf{x}}(T - s)\) satisfies: \[\partial_s \tilde p_{T-s}({\mathbf{x}}) = \nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, T - s) \tilde p_{T-s}({\mathbf{x}}) \bigr].\] Since \(\tilde{\mathbf{x}}(T) \sim p_T\), we have \(\tilde p_T = p_T\). The Fokker-Planck equation for \(p_t\) at \(t = T - s\) is: \[\partial_t p_{T-s}({\mathbf{x}}) = -\nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, T - s) p_{T-s}({\mathbf{x}}) \bigr].\] Since \(\partial_t = -\partial_s\), we get: \[\partial_s p_{T-s}({\mathbf{x}}) = \nabla_{{\mathbf{x}}} \cdot \bigl[ {\mathbf{v}}({\mathbf{x}}, T - s) p_{T-s}({\mathbf{x}}) \bigr].\] Both \(\tilde p_{T-s}\) and \(p_{T-s}\) satisfy the same PDE with the same initial condition at \(s = 0\) (\(\tilde p_T = p_T\)). Uniqueness implies \(\tilde p_{T-s} = p_{T-s}\), so \(\tilde{\mathbf{x}}(t) = \tilde{\mathbf{x}}(T - s) \sim p_{T-s} = p_t\), for all \(t \in [0, T]\).

Part 3: Reverse-Time SDE Marginal Densities.

We verify the marginal distributions using the forward-running reverse clock \[s:=T-t.\] Define \[{\mathbf{y}}(s):=\bar{\mathbf{x}}(T-s), \qquad s\in[0,T].\] Since the reverse process starts from \(\bar{\mathbf{x}}(T)\sim p_T\), we have \[{\mathbf{y}}(0)\sim p_T.\]

Under this change of clock, the reverse-time SDE becomes \[\mathop{}\!\mathrm{d}{\mathbf{y}}(s) = \Bigl[ -{\mathbf{f}}({\mathbf{y}}(s),T-s) + g^2(T-s)\nabla_{{\mathbf{x}}}\log p_{T-s}({\mathbf{y}}(s)) \Bigr]\mathop{}\!\mathrm{d}s + g(T-s)\mathop{}\!\mathrm{d}{\mathbf{w}}^{\mathrm{rev}}(s),\] where \({\mathbf{w}}^{\mathrm{rev}}\) is a standard Wiener process under the forward-running clock \(s\).

Let \(q_s\) denote the density of \({\mathbf{y}}(s)\). Its Fokker–Planck equation is \[\begin{align*} \partial_s q_s({\mathbf{x}}) &= -\nabla_{{\mathbf{x}}}\cdot \Bigl( \bigl[ -{\mathbf{f}}({\mathbf{x}},T-s) + g^2(T-s)\nabla_{{\mathbf{x}}}\log p_{T-s}({\mathbf{x}}) \bigr] q_s({\mathbf{x}}) \Bigr) \\ &\qquad + \tfrac12 g^2(T-s)\Delta_{{\mathbf{x}}}q_s({\mathbf{x}}). \end{align*}\]

We now check that \[q_s({\mathbf{x}})=p_{T-s}({\mathbf{x}})\] satisfies this equation. Using \[\nabla_{{\mathbf{x}}}\log p_{T-s}({\mathbf{x}})\,p_{T-s}({\mathbf{x}}) = \nabla_{{\mathbf{x}}}p_{T-s}({\mathbf{x}}),\] the right-hand side becomes \[\begin{align*} &\nabla_{{\mathbf{x}}}\cdot \bigl[ {\mathbf{f}}({\mathbf{x}},T-s)p_{T-s}({\mathbf{x}}) \bigr] - g^2(T-s)\Delta_{{\mathbf{x}}}p_{T-s}({\mathbf{x}}) \\ &\qquad + \tfrac12 g^2(T-s)\Delta_{{\mathbf{x}}}p_{T-s}({\mathbf{x}}) \\ &= \nabla_{{\mathbf{x}}}\cdot \bigl[ {\mathbf{f}}({\mathbf{x}},T-s)p_{T-s}({\mathbf{x}}) \bigr] - \tfrac12 g^2(T-s)\Delta_{{\mathbf{x}}}p_{T-s}({\mathbf{x}}). \end{align*}\]

On the other hand, the forward Fokker–Planck equation gives \[\partial_t p_t({\mathbf{x}}) = -\nabla_{{\mathbf{x}}}\cdot \bigl[{\mathbf{f}}({\mathbf{x}},t)p_t({\mathbf{x}})\bigr] + \tfrac12g^2(t)\Delta_{{\mathbf{x}}}p_t({\mathbf{x}}).\] Since \(t=T-s\), the chain rule gives \[\begin{align*} \partial_s p_{T-s}({\mathbf{x}}) &= -\partial_t p_t({\mathbf{x}})\big|_{t=T-s} \\ &= \nabla_{{\mathbf{x}}}\cdot \bigl[ {\mathbf{f}}({\mathbf{x}},T-s)p_{T-s}({\mathbf{x}}) \bigr] - \tfrac12g^2(T-s)\Delta_{{\mathbf{x}}}p_{T-s}({\mathbf{x}}). \end{align*}\] This is exactly the Fokker–Planck equation of the reversed process.

Finally, \(q_0=p_T\), which matches the initial distribution \({\mathbf{y}}(0)\sim p_T\). Under the regularity assumptions ensuring uniqueness of the Fokker–Planck solution, \[q_s=p_{T-s}, \qquad s\in[0,T].\] Returning to the original diffusion clock \(t=T-s\), we obtain \[\bar{\mathbf{x}}(t)\sim p_t.\] Thus, when the reverse-time SDE is solved backward from \(t=T\) to \(t=0\), it follows the same marginal distributions as the forward SDE, but in the reverse order. \(\blacksquare\)

Proposition in Chapter 4 (Minimizer of DSM): Minimizer of SM and DSM

Proof.

To find the minimizer \(\mathbf{s}^*\), we first consider a fixed time \(t\) and analyze the inner expectation in the objective function: \[\begin{align*} \mathcal J(t, \bm{\phi}) &:= \mathbb{E}_{{\mathbf{x}}_0 \sim p_{\text{data}}} \mathbb{E}_{{\mathbf{x}}_t \sim p_t(\cdot|{\mathbf{x}}_0)} \left[\left\| \mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t) - \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t|{\mathbf{x}}_0) \right\|_2^2 \right]. \end{align*}\] For this expectation to be minimized, we need to find \(\mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t)\) that minimizes the expected squared error for each \({\mathbf{x}}_t\). We can rewrite this expectation using the joint distribution of \(X_0\) and \(X_t\): \[\begin{align*} \mathcal J(t, \bm{\phi}) = \iint p_{\text{data}}({\mathbf{x}}_0) p_t({\mathbf{x}}_t |{\mathbf{x}}_0) \left\| \mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t) - \nabla_{{\mathbf{x}}_t} \log p({\mathbf{x}}_t | {\mathbf{x}}_0) \right\|_2^2 \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \mathop{}\!\mathrm{d}{\mathbf{x}}_t. \end{align*}\] For each fixed \({\mathbf{x}}_t\), we need to minimize: \[\begin{align*} \int p({\mathbf{x}}_0 | X_t = {\mathbf{x}}_t) p_t( {\mathbf{x}}_t) \left\| \mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t) - \nabla_{{\mathbf{x}}_t} \log p({\mathbf{x}}_t | {\mathbf{x}}_0) \right\|_2^2 \mathop{}\!\mathrm{d}{\mathbf{x}}_0. \end{align*}\] Since \(p_t({\mathbf{x}}_t)\) is constant with respect to \(\mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t)\), this is equivalent to minimizing: \[\begin{align*} \int p( {\mathbf{x}}_0 | X_t = {\mathbf{x}}_t) \left\| \mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t) - \nabla_{{\mathbf{x}}_t} \log p({\mathbf{x}}_t | {\mathbf{x}}_0) \right\|_2^2 \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \end{align*}\] This is minimized when \(\mathbf{s}_{\bm{\phi}}({\mathbf{x}}_t, t)\) equals the conditional expectation: \[\begin{align*} \mathbf{s}^*({\mathbf{x}}_t, t) = \mathbb{E}_{X_0 \sim p(X_0 | X_t = {\mathbf{x}}_t)} \left[ \nabla_{{\mathbf{x}}_t} \log p({\mathbf{x}}_t | X_0) \right]. \end{align*}\] Now we need to prove that this equals \(\nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t)\). By Bayes’ rule and the definition of marginal probability: \[\begin{align*} p_t({\mathbf{x}}_t) = \int p_t({\mathbf{x}}_t | {\mathbf{x}}_0) p_{\text{data}}({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0. \end{align*}\] Taking the logarithm and then the gradient with respect to \({\mathbf{x}}_t\): \[\begin{align*} \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t) = \frac{\nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t)}{p_t({\mathbf{x}}_t)} = \frac{\nabla_{{\mathbf{x}}_t} \int p_t({\mathbf{x}}_t | {\mathbf{x}}_0) p_{\text{data}}({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0} {\int p_t({\mathbf{x}}_t | {\mathbf{x}}_0) p_{\text{data}}({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0}. \end{align*}\] Under suitable regularity conditions, we can exchange the gradient and integral: \[\begin{align*} \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t) = \frac{\int \nabla_{{\mathbf{x}}_t} p_t({\mathbf{x}}_t | {\mathbf{x}}_0) p_{\text{data}}({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0} {\int p_t({\mathbf{x}}_t | {\mathbf{x}}_0) p_{\text{data}}({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0}. \end{align*}\] Using \[\nabla_{{\mathbf{x}}_t}p_t({\mathbf{x}}_t|{\mathbf{x}}_0) = p_t({\mathbf{x}}_t|{\mathbf{x}}_0) \nabla_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t|{\mathbf{x}}_0),\] and Bayes’ rule, we obtain \[\begin{align*} \nabla_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t) &= \int p({\mathbf{x}}_0|{\mathbf{x}}_t) \nabla_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t|{\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0 \\ &= \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)} \left[ \nabla_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t|{\mathbf{x}}_0) \right] = \mathbf{s}^*({\mathbf{x}}_t,t), \end{align*}\] which completes the proof. \(\blacksquare\)

Closed-Form Score Function of a Gaussian

For future reference, we summarize the formula for the score of a general multivariate normal distribution as the following lemma:

Flow-Based Perspective

Lemma in Chapter 5 (Instantaneous Change of Variables): Instantaneous Change of Variables

Proof.

Approach 1: Change-of-Variables Formula. We denote \(p({\mathbf{x}}(t), t)\) by \(p_t({\mathbf{x}}_t)\). Starting from the ODE discretization \[{\mathbf{z}}_{t+\Delta t} = {\mathbf{z}}_t + \Delta t {\mathbf{F}}({\mathbf{z}}_t, t),\] the change-of-variables formula for normalizing flows (Equation in Chapter 5) gives \[\begin{align*} \log p_{t+\Delta t}({\mathbf{z}}_{t+\Delta t}) &= \log p_t({\mathbf{z}}_t) - \log \Big| \det\big({\mathbf{I}}+ \Delta t \nabla_{{\mathbf{z}}} {\mathbf{F}}({\mathbf{z}}_t, t)\big) \Big| \\ &= \log p_t({\mathbf{z}}_t) - \mathop{\mathrm{Tr}}\Big(\log({\mathbf{I}}+ \Delta t \nabla_{{\mathbf{z}}} {\mathbf{F}}({\mathbf{z}}_t, t)) \Big) \\ &= \log p_t({\mathbf{z}}_t) - \Delta t \mathop{\mathrm{Tr}}\big(\nabla_{{\mathbf{z}}} {\mathbf{F}}({\mathbf{z}}_t, t)\big) + \mathcal{O}(\Delta t^2), \end{align*}\] where we used \(\log \det = \mathop{\mathrm{Tr}}\log\) and the expansion for small \(\Delta t\). Taking the limit \(\Delta t \to 0\) yields the continuous-time differential equation for the log-density. Indeed, the same trick is applied in Equation in Chapter 5.

Approach 2: Continuity Equation. We can also leverage the continuity equation, which essentially serves as the change-of-variables formula: \[\partial_t p({\mathbf{z}},t) = - \nabla_{{\mathbf{z}}} \cdot \big({\mathbf{F}}({\mathbf{z}},t) p({\mathbf{z}},t)\big).\] Expanding the divergence, \[\partial_t p = - \big( (\nabla_{{\mathbf{z}}} \cdot {\mathbf{F}}) p + {\mathbf{F}}\cdot \nabla_{{\mathbf{z}}} p \big).\] Along trajectories \({\mathbf{z}}(t)\) satisfying \(\frac{\mathop{}\!\mathrm{d}{\mathbf{z}}}{\mathop{}\!\mathrm{d}t} = {\mathbf{F}}({\mathbf{z}}(t), t)\), the total time derivative is \[\begin{align*} \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} p({\mathbf{z}}(t), t) &= \nabla_{{\mathbf{z}}} p \cdot \frac{\mathop{}\!\mathrm{d}{\mathbf{z}}}{\mathop{}\!\mathrm{d}t} + \partial_t p \\ &= \nabla_{{\mathbf{z}}} p \cdot {\mathbf{F}}- \big( (\nabla_{{\mathbf{z}}} \cdot {\mathbf{F}}) p + {\mathbf{F}}\cdot \nabla_{{\mathbf{z}}} p \big) \\ &= - (\nabla_{{\mathbf{z}}} \cdot {\mathbf{F}}) p. \end{align*}\] Dividing by \(p({\mathbf{z}}(t), t)\), we conclude \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} \log p({\mathbf{z}}(t), t) = - \nabla_{{\mathbf{z}}} \cdot {\mathbf{F}}({\mathbf{z}}(t), t).\] \(\blacksquare\)

Theorem in Chapter 5 (Mass Conservation Criterion) Mass Conservation Criterion

Some Prerequisites: Flow Map and Flow-Induced Density.

For any initial position \(\mathbf{x}_0\in\mathbb{R}^D\), the flow map \(\bm{\Psi}_t:\mathbb{R}^D\to\mathbb{R}^D\) is the unique solution of the ODE \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} \bm{\Psi}_t(\mathbf{x}_0) = \mathbf{v}_t\bigl(\bm{\Psi}_t(\mathbf{x}_0)\bigr), \quad \bm{\Psi}_0(\mathbf{x}_0)=\mathbf{x}_0.\] Under our regularity assumptions, \(\bm{\Psi}_t\) is continuously differentiable in both \(t\) and \(\mathbf{x}_0\).

The flow‑induced density \(p_t^{\mathrm{fwd}}\) is the pushforward of the initial density \(p_0\) by \(\bm{\Psi}_t\): \[p_t^{\mathrm{fwd}}(\mathbf{x}) = p_0\bigl(\bm{\Psi}_t^{-1}(\mathbf{x})\bigr) \Bigl|\det \bigl(\nabla\bm{\Psi}_t^{-1}(\mathbf{x})\bigr)\Bigr|.\] This gives the density at \(\mathbf{x}\) and time \(t\) of particles that started from \(p_0=p_{\mathrm{data}}\) and evolved under the velocity field \(\mathbf{v}_t\).

Informal Proof:

Sufficient Condition: \(p_t^{\mathrm{fwd}} = p_t\) \(\Rightarrow\) Continuity Equation. In Section (Continuous-Time Limit: Continuity Equation) in Appendix B we obtained a strong solution of the continuity equation by taking the continuous time limit of the discrete change of variable formula. In that approach, the density \(p_t\) is assumed smooth enough that all derivatives exist classically and the PDE holds pointwise. Here, we offer a complementary derivation in the weak sense: the continuity equation is imposed only after integrating against arbitrary smooth test functions, which relaxes the regularity requirements on both \(p_t\) and the velocity field \(\mathbf{v}_t\). This weak formulation is not only more rigorous since it accommodates less regular solutions but is also the standard framework in PDE theory and numerical analysis.

For any smooth test function \(\varphi(\mathbf{x})\) with compact support, the pushforward property gives: \[\begin{align*} \int p_t^{\mathrm{fwd}}(\mathbf{x}) \varphi(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x} &= \int p_0(\bm{\Psi}_t^{-1}(\mathbf{x})) \left| \det\left( \nabla \bm{\Psi}_t^{-1}(\mathbf{x}) \right) \right| \varphi(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x}\\ &= \int p_0(\mathbf{y}) \varphi(\bm{\Psi}_t(\mathbf{y})) \mathop{}\!\mathrm{d}\mathbf{y}, \end{align*}\] where the second equality follows from the change of variables \(\mathbf{x} = \bm{\Psi}_t(\mathbf{y})\), with \(\mathop{}\!\mathrm{d}\mathbf{y} = \left| \det\left( \nabla \bm{\Psi}_t^{-1}(\mathbf{x}) \right) \right| \mathop{}\!\mathrm{d}\mathbf{x}\).

Differentiate both sides with respect to \(t\): \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} \int p_t^{\mathrm{fwd}}(\mathbf{x}) \varphi(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x} = \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t} \int p_0(\mathbf{y}) \varphi(\bm{\Psi}_t(\mathbf{y})) \mathop{}\!\mathrm{d}\mathbf{y}.\]

The left-hand side is: \[\int \frac{\partial p_t^{\mathrm{fwd}}}{\partial t}(\mathbf{x}) \varphi(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x}.\]

On the right-hand side: \[\int p_0(\mathbf{y}) \nabla \varphi(\bm{\Psi}_t(\mathbf{y})) \cdot \mathbf{v}_t(\bm{\Psi}_t(\mathbf{y})) \mathop{}\!\mathrm{d}\mathbf{y},\] since \(\frac{\partial \bm{\Psi}_t}{\partial t}(\mathbf{y}) = \mathbf{v}_t(\bm{\Psi}_t(\mathbf{y}))\). Change variables to \(\mathbf{x} = \bm{\Psi}_t(\mathbf{y})\), so \[\mathop{}\!\mathrm{d}\mathbf{y} = \left| \det\left( \nabla \bm{\Psi}_t^{-1}(\mathbf{x}) \right) \right| \mathop{}\!\mathrm{d}\mathbf{x} ,\] and \[p_0(\mathbf{y}) = p_t^{\mathrm{fwd}}(\mathbf{x}) \left| \det\left( \nabla \bm{\Psi}_t(\mathbf{y}) \right) \right| = \frac{p_t^{\mathrm{fwd}}(\mathbf{x})}{\left| \det\left( \nabla \bm{\Psi}_t^{-1}(\mathbf{x}) \right) \right|}.\] Thus, the right-hand side becomes: \[\int p_t^{\mathrm{fwd}}(\mathbf{x}) \nabla \varphi(\mathbf{x}) \cdot \mathbf{v}_t(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x}.\]

Apply integration by parts, using the compact support of \(\varphi\): \[\int p_t^{\mathrm{fwd}}(\mathbf{x}) \nabla \varphi(\mathbf{x}) \cdot \mathbf{v}_t(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x} = -\int \varphi(\mathbf{x}) \nabla \cdot (p_t^{\mathrm{fwd}}(\mathbf{x}) \mathbf{v}_t(\mathbf{x})) \mathop{}\!\mathrm{d}\mathbf{x}.\] Equating both sides: \[\int \left[ \frac{\partial p_t^{\mathrm{fwd}}}{\partial t}(\mathbf{x}) + \nabla \cdot (p_t^{\mathrm{fwd}}(\mathbf{x}) \mathbf{v}_t(\mathbf{x})) \right] \varphi(\mathbf{x}) \mathop{}\!\mathrm{d}\mathbf{x} = 0.\] Since \(\varphi\) is arbitrary, we conclude: \[\frac{\partial p_t^{\mathrm{fwd}}}{\partial t} + \nabla \cdot (p_t^{\mathrm{fwd}} \mathbf{v}_t) = 0.\] Given \(p_t^{\mathrm{fwd}} = p_t\), this implies: \[\frac{\partial p_t}{\partial t} + \nabla \cdot (p_t \mathbf{v}_t) = 0.\]

Necessary Condition: Continuity Equation \(\Rightarrow p_t^{\mathrm{fwd}} = p_t\). Suppose \(p_t\) satisfies the continuity equation: \[\frac{\partial p_t}{\partial t} + \nabla \cdot (p_t \mathbf{v}_t) = 0,\] with the initial condition \(p_0(\mathbf{x}) = p_{\mathrm{data}}(\mathbf{x})\). We know \(p_t^{\mathrm{fwd}}\) satisfies the same continuity equation, as shown above, and: \[p_0^{\mathrm{fwd}}(\mathbf{x}) = p_0(\bm{\Psi}_0^{-1}(\mathbf{x})) \left| \det\left( \nabla \bm{\Psi}_0^{-1}(\mathbf{x}) \right) \right| = p_0(\mathbf{x}),\] since \(\bm{\Psi}_0(\mathbf{x}) = \mathbf{x}\). Thus, both densities share the same initial condition \(p_0 = p_{\mathrm{data}}\).

The continuity equation can be rewritten as: \[\frac{\partial p}{\partial t} + \mathbf{v}_t \cdot \nabla p + p \nabla \cdot \mathbf{v}_t = 0.\] This is a first-order linear PDE. Assuming \(\mathbf{v}_t\) is continuously differentiable and globally Lipschitz, and \(p_t\) is sufficiently smooth, the method of characteristics guarantees a unique solution in the space of smooth functions. Since \(p_t\) and \(p_t^{\mathrm{fwd}}\) satisfy the same PDE and initial condition, we conclude: \[p_t(\mathbf{x}) = p_t^{\mathrm{fwd}}(\mathbf{x})\] for all \(t\in[0,1]\) and \(\mathbf{x} \in \mathbb{R}^D\).

This completes the proof of the equivalence. \(\blacksquare\)

Theoretical Supplement: A Unified and Systematic View on Diffusion Models

Proposition in Chapter 6 (Equivalence of Parametrizations): Equivalence of Parametrizations

Proof:

As in Proposition in Chapter 4 (Minimizer of DSM) on the DSM loss, the global optimum of the matching loss \[\mathbb{E}_t\left[\omega(t)\,\mathbb{E}_{{\mathbf{x}}_0,\bm{\epsilon}}\left[\|\cdot\|_2^2\right]\right]\] is attained when the inner expectation \[\mathbb{E}_{{\mathbf{x}}_0,\bm{\epsilon}}\left[\|\cdot\|_2^2\right]\] is minimized for each fixed \(t\). Since this is a standard mean squared error problem, the minimizer is unique. From denoising score matching (Vincent 2011), Proposition in Chapter 4 (Minimizer of DSM) shows the optimal score function satisfies \[\begin{align*} {\mathbf{s}}^*({\mathbf{x}}_t,t) = \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}\left[\nabla_{{\mathbf{x}}}\log p_t({\mathbf{x}}_t|{\mathbf{x}}_0)\right] = \nabla_{{\mathbf{x}}_t} \log p_t({\mathbf{x}}_t). \end{align*}\] Using the identity \(\nabla_{{\mathbf{x}}}\log p_t({\mathbf{x}}_t|{\mathbf{x}}_0) = -\frac{1}{\sigma_t} \bm{\epsilon}\) for \(\bm{\epsilon} \sim \mathcal{N}(\bm{0},{\mathbf{I}})\), we obtain \[\begin{align*} {\mathbf{s}}^*({\mathbf{x}}_t,t) = -\frac{1}{\sigma_t}\, \bm{\epsilon}^*({\mathbf{x}}_t,t), \end{align*}\] where \(\bm{\epsilon}^*({\mathbf{x}}_t,t) = \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}[\bm{\epsilon}|{\mathbf{x}}_t]\) is the optimal \({\bm{\epsilon}}\)-prediction for \(\mathcal{L}_{\text{noise}}({\bm{\phi}})\). For the \({\mathbf{x}}\)-prediction loss \(\mathcal{L}_{\text{clean}}\), the optimal estimator is \[{\mathbf{x}}^*({\mathbf{x}}_t,t) = \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}[{\mathbf{x}}_0|{\mathbf{x}}_t],\] which, by Tweedie’s formula, relates to the score via \[\alpha_t\, {\mathbf{x}}^*({\mathbf{x}}_t,t) = {\mathbf{x}}_t + \sigma_t^2\, {\mathbf{s}}^*({\mathbf{x}}_t,t).\] For the \({\mathbf{v}}\)-prediction loss \(\mathcal{L}_{\text{velocity}}\), the optimal estimator is \[\begin{align*} {\mathbf{v}}^*({\mathbf{x}}_t,t) = \mathbb{E}_{p({\mathbf{x}}_0|{\mathbf{x}}_t)}[\alpha_t' {\mathbf{x}}_0 + \sigma_t' \bm{\epsilon}|{\mathbf{x}}_t] = \alpha_t' {\mathbf{x}}^* + \sigma_t' \bm{\epsilon}^*. \end{align*}\] Substituting the expressions for \({\mathbf{x}}^*\) and \(\bm{\epsilon}^*\) in terms of \({\mathbf{s}}^*\) yields \[\begin{align*} {\mathbf{v}}^*({\mathbf{x}}_t,t) = \frac{\alpha_t'}{\alpha_t} {\mathbf{x}}_t + \left( \frac{\alpha_t'}{\alpha_t} \sigma_t^2 - \sigma_t \sigma_t' \right) {\mathbf{s}}^*({\mathbf{x}}_t,t) = f(t) {\mathbf{x}}_t - \frac{1}{2}g^2(t) {\mathbf{s}}^*({\mathbf{x}}_t,t), \end{align*}\] which completes the derivation.

\(\blacksquare\)

Theoretical Supplement: Learning Fast Diffusion-Based Generators

Knowledge Distillation Loss as an Instance of the General Framework Equation in Chapter 10

We begin with the oracle KD loss \[\mathcal{L}_{\mathrm{KD}}^{\text{oracle}}({\bm{\theta}}) =\mathbb{E}_{{\mathbf{x}}_T\sim p_T}\, \big\|{\mathbf{G}}_{\bm{\theta}}({\mathbf{x}}_T,T,0)-{\bm{\Psi}}_{T\to 0}({\mathbf{x}}_T)\big\|_2^2,\] where \(p_T = p_{\mathrm{prior}}\). For the deterministic ODE flow map \({\bm{\Psi}}\) (which satisfies the semigroup property and is bijective along trajectories), the marginals satisfy the pushforward identities \[p_t = {\bm{\Psi}}_{0\to t}\,\sharp\,p_{\mathrm{data}} = {\bm{\Psi}}_{T\to t}\,\sharp\,p_{\mathrm{prior}}.\] Hence, \[{\bm{\Psi}}_{s\to T}\,\sharp\,p_s=p_T \quad\text{and}\quad {\bm{\Psi}}_{T\to 0}\circ{\bm{\Psi}}_{s\to T}={\bm{\Psi}}_{s\to 0}.\] Changing variables \({\mathbf{x}}_T={\bm{\Psi}}_{s\to T}({\mathbf{x}}_s)\) with \({\mathbf{x}}_s\sim p_s\) gives \[\mathcal{L}_{\mathrm{KD}}^{\text{oracle}}({\bm{\theta}}) = \mathbb{E}_{{\mathbf{x}}_s\sim p_s} \Big\|{\mathbf{G}}_{\bm{\theta}}\big({\bm{\Psi}}_{s\to T}({\mathbf{x}}_s),T,0\big) - {\bm{\Psi}}_{s\to 0}({\mathbf{x}}_s)\Big\|_2^2.\] Define the pulled-back student \(\widetilde{{\mathbf{G}}}_{\bm{\theta}}({\mathbf{x}}_s,s,0):= {\mathbf{G}}_{\bm{\theta}}({\bm{\Psi}}_{s\to T}({\mathbf{x}}_s),T,0)\) to express the same loss in the unified flow-map form (at \(t=0\)): \[\begin{align*} \mathcal{L}_{\mathrm{KD}}^{\text{oracle}}({\bm{\theta}}) = \mathbb{E}_{{\mathbf{x}}_s\sim p_s} \big\|\widetilde{{\mathbf{G}}}_{\bm{\theta}}({\mathbf{x}}_s,s,0)-{\bm{\Psi}}_{s\to 0}({\mathbf{x}}_s)\big\|_2^2. \end{align*}\] This derivation relies on change of variables through the oracle flow and the semigroup property.

\(\blacksquare\)

Transport Interpretation.

Proposition in Chapter 10 (θ-Gradient of L_VSD) also admits a distributional interpretation. For \[\hat{\mathbf{x}}_t = \alpha_t{\mathbf{G}}_{\bm{\theta}}({\mathbf{z}})+\sigma_t{\bm{\epsilon}},\] an infinitesimal parameter perturbation induces the data-space displacement \((\partial_{\bm{\theta}}\hat{\mathbf{x}}_t)\,\mathop{}\!\mathrm{d}{\bm{\theta}}\). Accordingly, the VSD gradient can be written as \[\nabla_{{\bm{\theta}}}\mathcal{L}_{\mathrm{VSD}} = \mathbb{E}\!\left[ \omega(t) \big(\partial_{\bm{\theta}}\hat{\mathbf{x}}_t\big)^\top \left( \nabla_{{\mathbf{x}}}\log p_t^{{\bm{\theta}}}(\hat{\mathbf{x}}_t) - \nabla_{{\mathbf{x}}}\log p_t(\hat{\mathbf{x}}_t) \right) \right].\] Thus, the score mismatch provides the data-space gradient signal, while the generator Jacobian maps this signal back to parameter space; gradient descent then updates the parameters in the opposite direction. This is the parametric counterpart of the distribution-level transport view in Section ((Optional) Wasserstein Gradient Flows as Distribution-Level Training) in Appendix B.

Theorem in Chapter 11 (CM-CT Equivalence up to Error O(Δ s^2)) CM Equals CT up to \(\mathcal{O}(\Delta s^2)\)

Proof:

Step 1: DDIM Update (with Oracle Score) Is the Conditional Mean. \[\begin{align*} \hat{\mathbf{x}}_{s'} &:= \frac{\alpha_{s'}}{\alpha_s} {\mathbf{x}}_s + \sigma_s^2 \left(\frac{\alpha_{s'}}{\alpha_s}-\frac{\sigma_{s'}}{\sigma_s}\right)\nabla_{{\mathbf{x}}_s}\log p_s({\mathbf{x}}_s) \\ &= \frac{\alpha_{s'}}{\alpha_s}\bigl({\mathbf{x}}_s+\sigma_s^2\nabla_{{\mathbf{x}}_s}\log p_s({\mathbf{x}}_s)\bigr) - \sigma_{s'}\sigma_s \nabla_{{\mathbf{x}}_s}\log p_s({\mathbf{x}}_s) \\ &= \alpha_{s'} \mathbb{E}[{\mathbf{x}}_0|{\mathbf{x}}_s] + \sigma_{s'}\Bigl(\frac{{\mathbf{x}}_s-\alpha_s \mathbb{E}[{\mathbf{x}}_0|{\mathbf{x}}_s]}{\sigma_s}\Bigr) \\ &= \alpha_{s'} \mathbb{E}[{\mathbf{x}}_0|{\mathbf{x}}_s] + \sigma_{s'} \mathbb{E}[{\bm{\epsilon}}|{\mathbf{x}}_s] \\ &= \mathbb{E}[\alpha_{s'}{\mathbf{x}}_0+\sigma_{s'}{\bm{\epsilon}}|{\mathbf{x}}_s] \\ &= \mathbb{E}[{\mathbf{x}}_{s'}|{\mathbf{x}}_s]. \end{align*}\] Here we use Tweedie’s formula \[\mathbb{E}[{\mathbf{x}}_0 |{\mathbf{x}}_s]=\frac{{\mathbf{x}}_s+\sigma_s^2\nabla_{{\mathbf{x}}_s}\log p_s({\mathbf{x}}_s)}{\alpha_s},\] together with the shared-noise coupling \[{\mathbf{x}}_s=\alpha_s{\mathbf{x}}_0+\sigma_s{\bm{\epsilon}}, \qquad {\mathbf{x}}_{s'}=\alpha_{s'}{\mathbf{x}}_0+\sigma_{s'}{\bm{\epsilon}} \quad (s'=s-\Delta s),\] in the third and fourth equalities. We also note that, more precisely, \(\hat{\mathbf{x}}_{s'}=\mathbb{E}[{\mathbf{x}}_{s'}|{\mathbf{x}}_s]\) should be read as \(\hat{\mathbf{x}}_{s'}=\mathbb{E}[{\mathbf{x}}_{s'}|{\mathbf{x}}_s, s]\), since \(s'=s-\Delta s\) is deterministically fixed once \(s\) is sampled.

Step 2: Expand CT Around the Conditional Mean \(\hat{\mathbf{x}}_{s'}\).

Recall that \(s\) is sampled first and we set \(s':=s-\Delta s\) deterministically. \[\begin{align*} \mathcal{L}_{\mathrm{CT}} &= \mathbb{E}_{s,{\mathbf{x}}_s}\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s} \Big[w(s) d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s), {\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_{s'},s')\big)\Big] \\ &\overset{\text{(1)}}{=} \mathbb{E}_{s,{\mathbf{x}}_s}\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s} \Big[ w(s) d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s), {\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s')\big) \\[-2pt] &\qquad\qquad\qquad\qquad + w(s) \partial_2 d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s),{\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s')\big)\big[\partial_1{\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s')({\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'})\big] \\[-2pt] &\qquad\qquad\qquad\qquad + w(s) \mathcal O \big(\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2\big)\Big] \\ &\overset{\text{(2)}}{=} \mathbb{E}_{s,{\mathbf{x}}_s} \Big[w(s) d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s), {\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s')\big)\Big] \\ &\quad+ \mathbb{E}_{s,{\mathbf{x}}_s} \Big[w(s) \partial_2 d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s),{\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s')\big)\big[\partial_1{\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s') \mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s}({\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'})\big]\Big] \\ &\quad+ \mathbb{E}_{s,{\mathbf{x}}_s}\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s} \Big[w(s) \mathcal{O} \big(\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2\big)\Big] \\ &\overset{\text{(3)}}{=} \mathbb{E}_{s,{\mathbf{x}}_s} \Big[w(s) d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s), {\mathbf{f}}_{{\bm{\theta}}^-}(\hat{\mathbf{x}}_{s'},s')\big)\Big] + \mathbb{E}_{s,{\mathbf{x}}_s}\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s} \Big[w(s) \mathcal{O} \big(\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2\big)\Big] \\ &= \mathcal{L}_{\mathrm{CM}} + \mathcal{O} \big(\mathbb{E}\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2\big) \\ &\overset{\text{(4)}}{=} \mathcal{L}_{\mathrm{CM}} + \mathcal{O}(\Delta s^2). \end{align*}\] Here, (1) applies a second-order Taylor expansion of \[h({\mathbf{x}}') := d\big({\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s), {\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}',s')\big)\] around \(\hat{\mathbf{x}}_{s'}\) in its second argument. (2) uses the tower property (a.k.a. the law of total expectation): \[\mathbb{E}_{s,{\mathbf{x}}_s}\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s}[\cdot] =\mathbb{E}_{s,{\mathbf{x}}_s}[\cdot],\] and notes that inside \(\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s}\), the terms \[w(s),\ {\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s),\ \hat{\mathbf{x}}_{s'},\ s'\] do not depend on the random variable \({\mathbf{x}}_{s'}\) (since \(s' = s-\Delta s\) and \(\hat{\mathbf{x}}_{s'}=\mathbb{E}[{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s]\) are fixed once \((s,{\mathbf{x}}_s)\) is fixed). Hence the only \({\mathbf{x}}_{s'}\)-dependence is through \(({\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'})\), and the inner expectations reduce to \(\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s}({\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'})\) and \(\mathbb{E}_{{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s}\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2\). (3) uses \[\mathbb{E}[{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s]=\bm 0 \quad\text{because}\quad \hat{\mathbf{x}}_{s'}=\mathbb{E}[{\mathbf{x}}_{s'}|{\mathbf{x}}_s,s].\] For (4), note that \[\mathbb{E}\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2 =\mathbb{E}\,\mathrm{Var}({\mathbf{x}}_{s'}|{\mathbf{x}}_s,s) \le \mathbb{E}\|{\mathbf{x}}_{s'}-{\mathbf{x}}_s\|^2.\] Moreover, \[{\mathbf{x}}_{s'}-{\mathbf{x}}_s=(\alpha_{s'}-\alpha_s){\mathbf{x}}_0+(\sigma_{s'}-\sigma_s){\bm{\epsilon}},\] and for a smooth scheduler, \(\alpha_{s'}-\alpha_s=\mathcal{O}(\Delta s)\) and \(\sigma_{s'}-\sigma_s=\mathcal{O}(\Delta s)\), hence \[\mathbb{E}\|{\mathbf{x}}_{s'}-{\mathbf{x}}_s\|^2=\mathcal{O}(\Delta s^2),\] which yields \[\mathbb{E}\|{\mathbf{x}}_{s'}-\hat{\mathbf{x}}_{s'}\|^2=\mathcal{O}(\Delta s^2).\]

\(\blacksquare\)

Proposition in Chapter 11 (Continuous-Time Consistency Training): Continuous-Time Limit of the CT Gradient

Proof:

We simplify the notation by working with Equation in Chapter 11 in the form \[\begin{align*} \mathcal{L}_{\mathrm{CM}}^{\Delta s}({\bm{\theta}},{\bm{\theta}}^-) := \mathbb{E}\left[\omega(s)\big\| {\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s) - {\mathbf{f}}_{{\bm{\theta}}^-}\big({\bm{\Psi}}_{s\to s-\Delta s}({\mathbf{x}}_s),\, s-\Delta s\big) \big\|_2^2\right]. \end{align*}\] For notational simplicity, write \[\tilde{\mathbf{x}}_{s-\Delta s}:={\bm{\Psi}}_{s\to s-\Delta s}({\mathbf{x}}_s).\]

Define \[\delta {\mathbf{f}} := {\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s)-{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s).\] Then \[\begin{align*} \big\|{\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s, s)- {\mathbf{f}}_{{\bm{\theta}}^-}(\tilde{\mathbf{x}}_{s-\Delta s}, s-\Delta s)\big\|_2^2 = \Big\|\delta {\mathbf{f}} + \Big({\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s, s)-{\mathbf{f}}_{{\bm{\theta}}^-}(\tilde{\mathbf{x}}_{s-\Delta s}, s-\Delta s)\Big)\Big\|_2^2. \end{align*}\]

Applying a Taylor expansion at \(({\mathbf{x}}_s,s)\) gives \[\begin{align*} {\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) - {\mathbf{f}}_{{\bm{\theta}}^-}(\tilde{\mathbf{x}}_{s-\Delta s}, s-\Delta s) = \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s)\,\Delta s + \mathcal{O}(|\Delta s|^2), \end{align*}\] where we use the first-order expansion of the oracle flow map, \[{\mathbf{x}}_s - {\bm{\Psi}}_{s\to s-\Delta s}({\mathbf{x}}_s) = {\mathbf{v}}^*({\mathbf{x}}_s,s)\,\Delta s + \mathcal{O}(|\Delta s|^2),\] together with the total differentiation identity \[\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) = \partial_s {\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) + \big(\partial_{{\mathbf{x}}}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s)\big)\,{\mathbf{v}}^*({\mathbf{x}}_s,s).\]

Therefore, \[\begin{align*} \big\|{\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s, s)- {\mathbf{f}}_{{\bm{\theta}}^-}(\tilde{\mathbf{x}}_{s-\Delta s}, s-\Delta s)\big\|_2^2 = \|\delta {\mathbf{f}}\|_2^2 + 2\,\delta {\mathbf{f}}^\top \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s)\,\Delta s + \mathcal{O}(|\Delta s|^2). \end{align*}\]

Now differentiate with respect to \({\bm{\theta}}\). Since \({\bm{\theta}}^-\) is treated as fixed with respect to \(\nabla_{{\bm{\theta}}}\) (i.e., no gradient is taken through the stop-gradient target), there is no gradient through \({\mathbf{f}}_{{\bm{\theta}}^-}\). Hence \[\begin{align*} \nabla_{{\bm{\theta}}}\mathcal{L}_{\mathrm{CM}}^{\Delta s}({\bm{\theta}},{\bm{\theta}}^-) &= \nabla_{{\bm{\theta}}}\mathbb{E}\!\left[\omega(s)\|\delta {\mathbf{f}}\|_2^2\right]\\ &\qquad\qquad + 2\,\mathbb{E}\!\left[ \omega(s)\, \nabla_{{\bm{\theta}}}{\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s)^\top \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) \right]\Delta s + \mathcal{O}(|\Delta s|^2). \end{align*}\]

At the comparison point where the target network is the stop-gradient copy of the online network, i.e., \({\bm{\theta}}={\bm{\theta}}^-\), we have \[\delta {\mathbf{f}} = {\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s)-{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) = 0,\] and therefore \[\nabla_{{\bm{\theta}}}\mathbb{E}\!\left[\omega(s)\|\delta {\mathbf{f}}\|_2^2\right]=0.\] Thus, \[\begin{align*} \left. \nabla_{{\bm{\theta}}}\mathcal{L}_{\mathrm{CM}}^{\Delta s}({\bm{\theta}},{\bm{\theta}}^-) \right|_{{\bm{\theta}}={\bm{\theta}}^-} &= \left. 2\,\mathbb{E}\!\left[ \omega(s)\, \nabla_{{\bm{\theta}}}{\mathbf{f}}_{{\bm{\theta}}}({\mathbf{x}}_s,s)^\top \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) \right] \right|_{{\bm{\theta}}={\bm{\theta}}^-}\Delta s + \mathcal{O}(|\Delta s|^2). \end{align*}\]

Dividing by \(\Delta s\) and taking the limit yields \[\begin{align*} \left. \lim_{\Delta s\to 0}\frac{1}{\Delta s}\nabla_{{\bm{\theta}}}\mathcal{L}_{\mathrm{CM}}^{\Delta s}({\bm{\theta}},{\bm{\theta}}^-) \right|_{{\bm{\theta}}={\bm{\theta}}^-} &= \left. \nabla_{{\bm{\theta}}}\mathbb{E}\!\left[ 2\,\omega(s)\, {\mathbf{f}}_{{\bm{\theta}}}^\top({\mathbf{x}}_s,s)\, \frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}s}{\mathbf{f}}_{{\bm{\theta}}^-}({\mathbf{x}}_s,s) \right] \right|_{{\bm{\theta}}={\bm{\theta}}^-}. \end{align*}\] This proves the identity at \({\bm{\theta}}={\bm{\theta}}^-\).

\(\blacksquare\)

(Optional) Elucidating Diffusion Model (EDM)

We introduce specific criteria for neural network parameterization design in the \({\mathbf{x}}\)-prediction model, as proposed in Elucidating Diffusion Models (EDM) (Karras et al. 2022). EDM provides a simple recipe that makes the training process easier to optimize and more reliable. The \({\mathbf{x}}\)-prediction model is expressed as a time-dependent skip connection combined with a scaled residual (Equation \(\eqref{eq:D-parametrization}\)). The central idea is to normalize both the inputs and the regression targets to unit variance at all times, and to adjust the skip path so that residual errors are not amplified as time evolves. This recipe has become a widely adopted default in diffusion model implementations and extends naturally to flow map learning, especially the family of Consistency Models.

Criteria for Neural Network \({\mathbf{x}}_{\bm{\phi}}\) Design

EDM considers a parametrization of the \({\mathbf{x}}\)-prediction model, denoted with a slight abuse of notation as \({\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}, t)\)¹, in the following form: \[\begin{align} \label{eq:D-parametrization} {\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}, t) := c_{\mathrm{skip}}(t) {\mathbf{x}}+ c_{\mathrm{out}}(t) {\mathbf{F}}_{\bm{\phi}}\left(c_{\mathrm{in}}(t) {\mathbf{x}}, c_{\mathrm{noise}}(t)\right). \end{align}\] Here, \(c_{\mathrm{skip}}(t)\), \(c_{\mathrm{out}}(t)\), \(c_{\mathrm{in}}(t)\), and \(c_{\mathrm{noise}}(t)\) are time-dependent functions. They are chosen to enhance stability and performance during training, based on practical considerations that will be introduced shortly.

Plugging this in Equation in Chapter 6, then the objective function becomes after straightforward algebraic manipulation²: \[\begin{align} \label{eq:loss-F} \mathbb{E}_{{\mathbf{x}}_0, \bm{\epsilon}, t} \left[\omega(t)c_{\mathrm{out}}^2(t)\left\lVert {\mathbf{F}}_{\bm{\phi}}\left(c_{\mathrm{in}}(t){\mathbf{x}}_t, c_{\mathrm{noise}}(t)\right) - {\mathbf{x}}_{\mathrm{tgt}}(t)\right\rVert_2^2\right]. \end{align}\] Here, the regression target \({\mathbf{x}}_{\mathrm{tgt}}(t)\) is obtained as: \[{\mathbf{x}}_{\mathrm{tgt}}(t) = \frac{\left(1-c_{\mathrm{skip}}(t)\alpha_t\right){\mathbf{x}}_0-c_{\mathrm{skip}}(t)\sigma_t \bm{\epsilon}}{c_{\mathrm{out}}(t)}.\]

By incorporating the standard deviation of \(p_{\text{data}}\), denoted as \(\sigma_{\mathrm{d}}\), EDM proposes the following design criterion for network parameterization, which may be described as the unit variance criterion.

Unit Variance of Input.

\[\begin{align*} &\text{Var}_{{\mathbf{x}}_0,\bm{\epsilon}}\left[ c_{\mathrm{in}}(t){\mathbf{x}}_t\right] = 1 \nonumber \\ \Longleftrightarrow &~c_{\mathrm{in}}^2(t) \text{Var}_{{\mathbf{x}}_0,\bm{\epsilon}}\left[ \alpha_t {\mathbf{x}}_0 +\sigma_t \bm{\epsilon} \right] = 1 \nonumber \\ \Longleftrightarrow &~c_{\mathrm{in}}(t) = \frac{1}{\sqrt{\sigma_{\mathrm{d}}^2 \alpha_t^2 + \sigma_t^2}}, \end{align*}\] taking the positive root.

Unit Variance of Training Target.

\[\begin{align} \label{eq:c-out} &\text{Var}_{{\mathbf{x}}_0,\bm{\epsilon}}\left[ {\mathbf{x}}_{\mathrm{tgt}}(t) \right] = 1 \nonumber \\ \Longleftrightarrow &~c_{\mathrm{out}}^2(t) = \left(1-c_{\mathrm{skip}}(t)\alpha_t\right)^2 \sigma_{\mathrm{d}}^2 + c_{\mathrm{skip}}^2(t)\sigma_t^2, \end{align}\] with centered data (\(\mathbb{E}[{\mathbf{x}}_0]=\bm{0}\)).

Minimize the Error Amplification from \({\mathbf{F}}_{\bm{\phi}}\) to \({\mathbf{x}}_{\bm{\phi}}\).

EDM aims to mitigate the amplification of the network’s learning error from \({\mathbf{F}}_{\bm{\phi}}\) to \({\mathbf{x}}_{\bm{\phi}}\). This is achieved by selecting \(c_{\mathrm{skip}}\) to minimize \(c_{\mathrm{out}}\): \[c_{\mathrm{skip}}^* \in \mathop{\mathrm{arg\,min}}_{c_{\mathrm{skip}}} c_{\mathrm{out}}^2.\] Using the standard approach of solving \(\frac{\partial c_{\mathrm{out}}}{\partial c_{\mathrm{skip}}} = 0\) for the critical point \(c_{\mathrm{skip}}^*\), we obtain \[\begin{align*} c_{\mathrm{skip}}^*(t) = \frac{\alpha_t \sigma_{\mathrm{d}}^2 }{\alpha_t^2 \sigma_{\mathrm{d}}^2 + \sigma_t^2}. \end{align*}\] Substituting this into Equation \(\eqref{eq:c-out}\), the optimal value is given by \[\begin{align*} c_{\mathrm{out}}^*(t) = \pm \frac{\sigma_t \sigma_{\mathrm{d}} }{\sqrt{\alpha_t^2 \sigma_{\mathrm{d}}^2 + \sigma_t^2}}. \end{align*}\] By convention, we use the nonnegative branch for the output scale: \[c_{\mathrm{out}}^*(t) = \frac{\sigma_t \sigma_{\text d}} {\sqrt{\alpha_t^2 \sigma_{\text d}^2 + \sigma_t^2}} \;\;\;(\ge 0),\] which ensures \(c_{\mathrm{out}}(0)=0\) and \(c_{\mathrm{out}}(t)\to\sigma_{\text d}\) as \(\sigma_t\) is sufficiently large, yielding the intuitive limits \({\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}_t,0)\approx {\mathbf{x}}_t\) and \({\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}_t,t)\approx \sigma_{\text d}\,{\mathbf{F}}_{\bm{\phi}}(\cdot)\).

We summarize these coefficients as follows:

With Equation \(\eqref{eq:edm-coeff}\), the regression target \({\mathbf{x}}_{\mathrm{tgt}}(t)\) simplifies to \[{\mathbf{x}}_{\mathrm{tgt}}(t) = \frac{1}{\sigma_{\mathrm{d}}}\frac{\sigma_t {\mathbf{x}}_0 - \alpha_t \sigma_{\mathrm{d}}^2\bm{\epsilon} }{\sqrt{R_t}}.\] Additionally, substituting these expressions into Equation \(\eqref{eq:D-parametrization}\) and Equation \(\eqref{eq:loss-F}\) allows us to simplify both the parametrization and the loss function, yielding: \[\begin{align*} {\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}, t)=\frac{\alpha_t \sigma_{\mathrm{d}}^2 }{R_t} {\mathbf{x}}+ \frac{\sigma_t \sigma_{\mathrm{d}} }{\sqrt{R_t}}{\mathbf{F}}_{\bm{\phi}}\left(\frac{1}{\sqrt{R_t}}{\mathbf{x}}, c_{\mathrm{noise}}(t)\right), \end{align*}\] and \[\begin{align*} \mathbb{E}_{{\mathbf{x}}_0, \bm{\epsilon}, t} \left[\omega(t)\frac{\sigma_t^2 }{R_t}\left\lVert \sigma_{\mathrm{d}} {\mathbf{F}}_{\bm{\phi}}\left(\frac{1}{\sqrt{R_t}}{\mathbf{x}}_t, c_{\mathrm{noise}}(t)\right) - \left(\frac{\sigma_t {\mathbf{x}}_0 - \alpha_t \sigma_{\mathrm{d}}^2\bm{\epsilon} }{\sqrt{R_t}}\right)\right\rVert_2^2\right]. \end{align*}\]

With this parameterization and the conditions \(\alpha_0 \approx 1\) and \(\sigma_0 \approx 0\), we observe that \[c_{\mathrm{skip}}(0) \approx 1 \quad \text{and} \quad c_{\mathrm{out}}(0) \approx 0.\]

Selection of \(c_{\mathrm{noise}}(t)\).

It provides a noise-level embedding to \({\mathbf{F}}_{\bm{\phi}}\); any one-to-one mapping of the noise level \(\sigma_t\) (e.g., \(c_{\mathrm{noise}}(t)=\log \sigma_t\)) is suitable.

A Common EDM Special Case: \(\alpha_t = 1\), \(\sigma_t = t\)

We consider the simplified noise schedule used in EDM, where \(\alpha_t = 1\) and \(\sigma_t = t\), which also appears in the discussion of CTM in Section (General Flow Map: Consistency Trajectory ModelGeneral Flow Map: Consistency Trajectory Model) in Chapter 11. Under this setting, the forward process becomes \[{\mathbf{x}}_t = {\mathbf{x}}_0 + t\bm{\epsilon}, \quad \text{with } {\mathbf{x}}_0 \sim p_{\text{data}}, \,\, \bm{\epsilon} \sim \mathcal{N}(\bm{0}, {\mathbf{I}}),\] corresponding to the perturbation kernel \[p_t({\mathbf{x}}_t|{\mathbf{x}}_0) = \mathcal{N}({\mathbf{x}}_t; {\mathbf{x}}_0, t^2{\mathbf{I}}).\]

The resulting marginal density is \[p_t({\mathbf{x}}) = \int \mathcal{N}({\mathbf{x}};{\mathbf{x}}_0,t^2{\mathbf{I}}) p_{\mathrm{data}}({\mathbf{x}}_0) \mathop{}\!\mathrm{d}{\mathbf{x}}_0 = \bigl[ p_{\mathrm{data}} * \mathcal{N}(\bm{0},t^2{\mathbf{I}}) \bigr]({\mathbf{x}}).\] Thus, the exact terminal marginal is \(p_T=p_{\mathrm{data}}*\mathcal{N}(\bm{0},T^2{\mathbf{I}})\). When \(T\) is sufficiently large relative to the scale of the data, this marginal is commonly approximated by \(\mathcal{N}(\bm{0},T^2{\mathbf{I}})\), which is used as the sampling prior: \[p_{\mathrm{prior}}({\mathbf{x}}_T) := \mathcal{N}({\mathbf{x}}_T;\bm{0},T^2{\mathbf{I}}).\]

Under this setup, the PF-ODE based on \({\mathbf{x}}\)-prediction \({\mathbf{x}}_{\bm{\phi}^\times}\) (see Equation in Chapter 6) simplifies to \[\begin{align*} \frac{\mathop{}\!\mathrm{d}{\mathbf{x}}(t)}{\mathop{}\!\mathrm{d}t} = \frac{{\mathbf{x}}(t) - {\mathbf{x}}_{\bm{\phi}^\times}({\mathbf{x}}(t), t)}{t}. \end{align*}\]

Substituting this formulation into Equation \(\eqref{eq:edm-coeff}\), the neural network parameterization coefficients become \[\begin{align} \label{eq:edm-simple-coefficients} c_{\mathrm{in}}(t) = \frac{1}{\sqrt{\sigma_{\mathrm{d}}^2 + t^2}}, \quad c_{\mathrm{skip}}(t) = \frac{\sigma_{\mathrm{d}}^2}{\sigma_{\mathrm{d}}^2 + t^2}, \quad c_{\mathrm{out}}(t) = \pm \frac{t \sigma_{\mathrm{d}}}{\sqrt{\sigma_{\mathrm{d}}^2 + t^2}}. \end{align}\] From these expressions, we observe:

• When \(t \approx 0\), the noise level is negligible, so \(c_{\mathrm{skip}} \approx 1\) and \(c_{\mathrm{out}} \approx 0\). In this limit, the skip path dominates and the network essentially passes through its input, \[{\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}, t) \approx {\mathbf{x}}.\]

• When \(t \gg 0\), the input is heavily corrupted by noise, so \(c_{\mathrm{skip}} \approx 0\) and \(c_{\mathrm{out}} \approx \sigma_{\mathrm{d}}\). In this regime, the skip path vanishes and the model output is determined entirely by the learned residual, \[{\mathbf{x}}_{\bm{\phi}}({\mathbf{x}}, t) \approx \sigma_{\mathrm{d}} {\mathbf{F}}_{\bm{\phi}}\left(c_{\mathrm{in}}(t){\mathbf{x}}, c_{\mathrm{noise}}(t)\right),\] meaning that the network \({\mathbf{F}}_{\bm{\phi}}\) predicts a scaled proxy of the clean signal from a normalized noisy input; at high noise levels the model output is therefore determined entirely by the learned denoising function.

In short, the parameterization smoothly interpolates from an identity map at small \(t\) to a scaled residual predictor on standardized inputs at large \(t\).