Rambling on About Generative Diffusion Models (Part 6): The ODE Formulation of the General Framework

生成扩散模型漫谈（六）：一般框架之ODE篇 苏剑林 · 2022-08-08 · Machine Learning differential equation generative model DDPM diffusion

In the previous post Rambling on About Generative Diffusion Models (Part 5): The SDE Formulation of the General Framework, we gave a basic introduction to and derivation of Dr. Yang Song's paper Score-Based Generative Modeling through Stochastic Differential Equations. However, as the title suggests, that previous post was mainly concerned with the SDE-related parts of the original paper, leaving aside the part known as the "Probability Flow ODE". This post fills in that gap.

In fact, this leftover content only occupies a small section in the body of the original paper, but I found, after mulling it over for quite a while, that we still need a whole new post to introduce it—because the derivation of this result cannot avoid the Fokker-Planck equation. So we'll need some space to introduce the Fokker-Planck equation before our star, the ODE, can take the stage.

Reflecting Again

Let's briefly recap the content of the previous post. First, we used an SDE to define a forward process (the "demolition" process):

\begin{equation}d\boldsymbol{x} = \boldsymbol{f}_t(\boldsymbol{x}) dt + g_t d\boldsymbol{w}\label{eq:sde-forward}\end{equation}more

Then, we derived the corresponding reverse-process SDE (the "construction" process):

\begin{equation}d\boldsymbol{x} = \left[\boldsymbol{f}_t(\boldsymbol{x}) - g_t^2\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}) \right] dt + g_t d\boldsymbol{w}\label{eq:sde-reverse}\end{equation}

Finally, we derived the loss function (score matching) for using a neural network $\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}, t)$ to estimate $\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})$:

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{x}_t \sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)}\left[\left\Vert \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\right] \end{equation}

At this point, we've completed the general framework for training and sampling in diffusion models—it can be said to be a very general generalization of DDPM. But just as Rambling on About Generative Diffusion Models (Part 4): DDIM = A High-Level Perspective on DDPM showed that DDIM is the result of reflecting on DDPM from a "higher vantage point", does the SDE—being a generalization of DDPM—also have a corresponding "high-level reflection result"? It does, and that result is the subject of this post: the "Probability Flow ODE".

The Dirac Function

What kind of reflection did DDIM carry out? Quite simply, DDIM discovered that DDPM's training objective mainly depends on $p(\boldsymbol{x}_t|\boldsymbol{x}_0)$, and is unrelated to $p(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})$. So it took $p(\boldsymbol{x}_t|\boldsymbol{x}_0)$ as the starting point and derived more general forms of $p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)$ and $p(\boldsymbol{x}_t|\boldsymbol{x}_{t-1},\boldsymbol{x}_0)$. The reflection carried out by the probability flow ODE is similar in spirit: it asks, within the SDE framework, for a fixed $p(\boldsymbol{x}_t)$, what different choices of $p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t)$ are possible (equivalently, what different forward-process SDEs can we find)?

Let's first write out the discretized form of the forward process $\eqref{eq:sde-forward}$:

\begin{equation}\boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\label{eq:sde-discrete}\end{equation}

This equation describes a relationship between the random variables $\boldsymbol{x}_{t+\Delta t},\boldsymbol{x}_t,\boldsymbol{\varepsilon}$, and we could conveniently take expectations of both sides. However, what we actually want is not an expectation but rather the relationship satisfied by the distribution $p(\boldsymbol{y}_t)$. How do we convert a distribution into an expectation form? The answer is the Dirac function:

\begin{equation}p(\boldsymbol{x}) = \int \delta(\boldsymbol{x} - \boldsymbol{y}) p(\boldsymbol{y}) d\boldsymbol{y} = \mathbb{E}_{\boldsymbol{y}}[\delta(\boldsymbol{x} - \boldsymbol{y})]\end{equation}

The rigorous definition of the Dirac function belongs to the realm of functional analysis, but we usually just treat it as an ordinary function, and this generally gives correct results. From the equation above, we can also deduce that, for any $f(\boldsymbol{x})$,

\begin{equation}p(\boldsymbol{x})f(\boldsymbol{x}) = \int \delta(\boldsymbol{x} - \boldsymbol{y}) p(\boldsymbol{y})f(\boldsymbol{y}) d\boldsymbol{y} = \mathbb{E}_{\boldsymbol{y}}[\delta(\boldsymbol{x} - \boldsymbol{y}) f(\boldsymbol{y})]\end{equation}

Taking the partial derivative of both sides directly, we get

\begin{equation}\nabla_{\boldsymbol{x}}[p(\boldsymbol{x}) f(\boldsymbol{x})] = \mathbb{E}_{\boldsymbol{y}}\left[\nabla_{\boldsymbol{x}}\delta(\boldsymbol{x} - \boldsymbol{y}) f(\boldsymbol{y})\right] = \mathbb{E}_{\boldsymbol{y}}\left[f(\boldsymbol{y})\nabla_{\boldsymbol{x}}\delta(\boldsymbol{x} - \boldsymbol{y})\right]\end{equation}

This is one of the properties we'll need later; essentially it shows that a derivative of the Dirac function can, via integration, be transferred onto the function it's multiplying.

The F-P Equation

With this groundwork laid, we can now use equation $\eqref{eq:sde-discrete}$ to write

\begin{equation}\begin{aligned} &\,\delta(\boldsymbol{x} - \boldsymbol{x}_{t+\Delta t}) \\[5pt] =&\, \delta(\boldsymbol{x} - \boldsymbol{x}_t - \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t - g_t \sqrt{\Delta t}\boldsymbol{\varepsilon}) \\ \approx&\, \delta(\boldsymbol{x} - \boldsymbol{x}_t) - \left(\boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon}\right)\cdot \nabla_{\boldsymbol{x}}\delta(\boldsymbol{x} - \boldsymbol{x}_t) + \frac{1}{2} \left(g_t\sqrt{\Delta t}\boldsymbol{\varepsilon}\cdot \nabla_{\boldsymbol{x}}\right)^2\delta(\boldsymbol{x} - \boldsymbol{x}_t) \end{aligned}\end{equation}

Here, we treated $\delta(\cdot)$ as an ordinary function and performed a Taylor expansion, keeping only terms up to order $\mathcal{O}(\Delta t)$. Now let's take the expectation of both sides:

\begin{equation}\begin{aligned} &\,p_{t+\Delta t}(\boldsymbol{x}) \\[6pt] =&\,\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\delta(\boldsymbol{x} - \boldsymbol{x}_{t+\Delta t})\right] \\ \approx&\, \mathbb{E}_{\boldsymbol{x}_t, \boldsymbol{\varepsilon}}\left[\delta(\boldsymbol{x} - \boldsymbol{x}_t) - \left(\boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon}\right)\cdot \nabla_{\boldsymbol{x}}\delta(\boldsymbol{x} - \boldsymbol{x}_t) + \frac{1}{2} \left(g_t\sqrt{\Delta t}\boldsymbol{\varepsilon}\cdot \nabla_{\boldsymbol{x}}\right)^2\delta(\boldsymbol{x} - \boldsymbol{x}_t)\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\delta(\boldsymbol{x} - \boldsymbol{x}_t) - \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t\cdot \nabla_{\boldsymbol{x}}\delta(\boldsymbol{x} - \boldsymbol{x}_t) + \frac{1}{2} g_t^2\Delta t \nabla_{\boldsymbol{x}}\cdot \nabla_{\boldsymbol{x}}\delta(\boldsymbol{x} - \boldsymbol{x}_t)\right] \\ =&\,p_t(\boldsymbol{x}) - \nabla_{\boldsymbol{x}}\cdot\left[\boldsymbol{f}_t(\boldsymbol{x})\Delta t\, p_t(\boldsymbol{x})\right] + \frac{1}{2}g_t^2\Delta t \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x}) \end{aligned}\end{equation}

Dividing both sides by $\Delta t$ and taking the limit $\Delta t\to 0$, we obtain

\begin{equation}\frac{\partial}{\partial t} p_t(\boldsymbol{x}) = - \nabla_{\boldsymbol{x}}\cdot\left[\boldsymbol{f}_t(\boldsymbol{x}) p_t(\boldsymbol{x})\right] + \frac{1}{2}g_t^2 \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x})\label{eq:fp} \end{equation}

This is the "F-P equation" (Fokker-Planck equation) corresponding to equation $\eqref{eq:sde-forward}$; it's a partial differential equation describing the marginal distribution.

An Equivalent Transformation

Don't worry about seeing a PDE here—we have no intention of studying how to solve it, and are only using it to guide an equivalent transformation. For any function $\sigma_t$ satisfying $\sigma_t^2\leq g_t^2$, the F-P equation $\eqref{eq:fp}$ is completely equivalent to

\begin{equation}\begin{aligned} \frac{\partial}{\partial t} p_t(\boldsymbol{x}) =&\, - \nabla_{\boldsymbol{x}}\cdot\left[\boldsymbol{f}_t(\boldsymbol{x})p_t(\boldsymbol{x}) - \frac{1}{2}(g_t^2 - \sigma_t^2)\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x})\right] + \frac{1}{2}\sigma_t^2 \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x}) \\ =&\,- \nabla_{\boldsymbol{x}}\cdot\left[\left(\boldsymbol{f}_t(\boldsymbol{x}) - \frac{1}{2}(g_t^2 - \sigma_t^2)\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\right)p_t(\boldsymbol{x})\right] + \frac{1}{2}\sigma_t^2 \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x}) \end{aligned}\label{eq:fp-2}\end{equation}

Formally, this F-P equation is the same as the original one but with $\boldsymbol{f}_t(\boldsymbol{x})$ replaced by $\boldsymbol{f}_t(\boldsymbol{x}) - \frac{1}{2}(g_t^2 - \sigma_t^2)\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})$ and $g_t$ replaced by $\sigma_t$. Since equation $\eqref{eq:fp}$ corresponds to equation $\eqref{eq:sde-forward}$, the above then corresponds to

\begin{equation}d\boldsymbol{x} = \left(\boldsymbol{f}_t(\boldsymbol{x}) - \frac{1}{2}(g_t^2 - \sigma_t^2)\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\right) dt + \sigma_t d\boldsymbol{w}\label{eq:sde-forward-2}\end{equation}

But don't forget that equation $\eqref{eq:fp}$ and equation $\eqref{eq:fp-2}$ are completely equivalent, which means the marginal distributions $p_t(\boldsymbol{x})$ corresponding to the two stochastic differential equations $\eqref{eq:sde-forward}$ and $\eqref{eq:sde-forward-2}$ are exactly the same! This result tells us that forward processes with different variances can produce identical marginal distributions. This result is essentially an upgraded version of DDIM—and later we'll show that when $\boldsymbol{f}_t(\boldsymbol{x})$ is a linear function of $\boldsymbol{x}$, this reduces exactly to DDIM.

In particular, based on the SDE results from the previous post, we can write down the reverse SDE corresponding to equation $\eqref{eq:sde-forward-2}$:

\begin{equation}d\boldsymbol{x} = \left(\boldsymbol{f}_t(\boldsymbol{x}) - \frac{1}{2}(g_t^2 + \sigma_t^2)\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\right) dt + \sigma_t d\boldsymbol{w}\label{eq:sde-reverse-2}\end{equation}

Neural ODE

Equation $\eqref{eq:sde-forward-2}$ allows us to change the variance of the sampling process. Here we're particularly interested in the extreme case where $\sigma_t = 0$, in which the SDE degenerates into an ODE (ordinary differential equation):

\begin{equation}d\boldsymbol{x} = \left(\boldsymbol{f}_t(\boldsymbol{x}) - \frac{1}{2}g_t^2\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\right) dt\label{eq:flow-ode}\end{equation}

This ODE is called the "Probability Flow ODE". Since in practice $\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})$ needs to be approximated by a neural network $\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}, t)$, the equation above also corresponds to a "Neural ODE".

Why is it worth specifically studying the zero-variance case? Because in this case the propagation process is noise-free—the transformation from $\boldsymbol{x}_0$ to $\boldsymbol{x}_T$ is deterministic, so by directly solving the ODE in reverse we can obtain the inverse transformation from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$, which is also a deterministic transformation. (Substituting $\sigma_t=0$ directly into equation $\eqref{eq:sde-reverse-2}$ also reveals that the forward and reverse equations are identical.) This process is consistent with flow models (i.e., transforming noise into samples via an invertible mapping), so the probability flow ODE lets us connect the results of diffusion models with those of flow models. For instance, the original paper mentions that the probability flow ODE allows for exact likelihood computation and obtaining latent representations—these are essentially benefits inherited from flow models. Because flow models are invertible, they also allow us to perform various image editing operations on the original image in latent space.

On another note, since the transformation from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$ is described by an ODE, this means we can speed up the transformation from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$ using various high-order numerical ODE solvers. Of course, in principle there are also acceleration methods for solving SDEs, but research on accelerating SDEs is far less developed and far less straightforward than that for ODEs. Overall, compared to SDEs, ODEs are simpler and more direct both for theoretical analysis and practical solving.

Revisiting DDIM

At the end of Rambling on About Generative Diffusion Models (Part 4): DDIM = A High-Level Perspective on DDPM, we derived that the continuous version of DDIM corresponds to the ODE

\begin{equation}\frac{d}{ds}\left(\frac{\boldsymbol{x}(s)}{\bar{\alpha}(s)}\right) = \boldsymbol{\epsilon}_{\boldsymbol{\theta}}\left(\boldsymbol{x}(s), t(s)\right)\frac{d}{ds}\left(\frac{\bar{\beta}(s)}{\bar{\alpha}(s)}\right)\label{eq:ddim-ode}\end{equation}

We can now see that this result is in fact just a special case of equation $\eqref{eq:flow-ode}$ in this post, when $\boldsymbol{f}_t(\boldsymbol{x})$ is taken to be the linear function $f_t \boldsymbol{x}$: at the end of Rambling on About Generative Diffusion Models (Part 5): The SDE Formulation of the General Framework, we derived the corresponding relations

\begin{equation}\left\{\begin{aligned} &f_t = \frac{1}{\bar{\alpha}_t}\frac{d\bar{\alpha}_t}{dt} \\ &g^2 (t) = 2\bar{\alpha}_t \bar{\beta}_t \frac{d}{dt}\left(\frac{\bar{\beta}_t}{\bar{\alpha}_t}\right) \\ &\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}, t) = -\frac{\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}, t)}{\bar{\beta}_t} \end{aligned}\right.\end{equation}

Substituting these relations into equation $\eqref{eq:flow-ode}$ [with $\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})$ replaced by $\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}, t)$] and simplifying, we get

\begin{equation}\frac{1}{\bar{\alpha}_t}\frac{d\boldsymbol{x}}{dt} - \frac{\boldsymbol{x}}{\bar{\alpha}_t^2}\frac{d\bar{\alpha}_t}{dt} = \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\frac{d}{dt}\left(\frac{\bar{\beta}_t}{\bar{\alpha}_t}\right)\end{equation}

The left-hand side can be further simplified to $\frac{d}{dt}\left(\frac{\boldsymbol{x}}{\bar{\alpha}_t}\right)$, so the equation above is exactly equivalent to equation $\eqref{eq:ddim-ode}$.

Summary

Building on the SDE post, this post used the F-P equation to derive a more general forward equation, from which we obtained the "Probability Flow ODE", and we showed that DDIM is a special case of it.