Rambling on About Generative Diffusion Models (12): Tackling the Diffusion ODE Head-On

生成扩散模型漫谈（十二）：“硬刚”扩散ODE 苏剑林 · 2022-09-28 · Machine Learning differential equation generative model diffusion

In Rambling on About Generative Diffusion Models (5): The General Framework in the SDE Perspective, we understood generative diffusion models from the perspective of the SDE, and then in Rambling on About Generative Diffusion Models (6): The General Framework in the ODE Perspective, we learned that the diffusion model corresponding to the SDE actually has an ODE model implicit within it. Interestingly, in Rambling on About Generative Diffusion Models (4): DDIM = DDPM from a Higher Vantage Point we also learned that the originally stochastic sampling process of DDPM has, hidden within it, a deterministic sampling procedure called DDIM, whose continuous limit is likewise an ODE.

Thinking this through, we can see that both "DDPM → DDIM" and "SDE → ODE" are transitions from a stochastic sampling model to a deterministic one. But if our end goal from the very start is an ODE, this route seems somewhat "roundabout." In this post, I'll try to give a direct derivation of the ODE-based diffusion model, and along the way reveal its connections to the Jacobian determinant, the heat equation, and related topics.

Differential Equations

Generative models like GANs are, at heart, aiming to find a deterministic transformation that maps random variables sampled from a simple distribution (such as a standard normal) into samples from a particular data distribution. Flow-based models are also a kind of generative model, but they take the opposite route: first find an invertible transformation that maps the data distribution to a simple distribution, and then solve for the corresponding inverse transformation to obtain a generative model.

Traditional flow models achieve this invertible transformation by designing elaborate coupling layers (see the series "Flow-Based Models, One Step at a Time"). But eventually people realized that this transformation can also be realized through differential equations, and in a theoretically quite elegant way. A whole line of research on building generative models via "neural networks + differential equations" forms a subfield now known as "Neural ODEs."

Consider a first-order (system of) ordinary differential equation(s) on $\boldsymbol{x}_t\in\mathbb{R}^d$:

\begin{equation}\frac{d\boldsymbol{x}_t}{dt}=\boldsymbol{f}_t(\boldsymbol{x}_t)\label{eq:ode}\end{equation}

Suppose $t\in[0, T]$. Then, given $\boldsymbol{x}_0$, (under fairly easily satisfied conditions) we can deterministically solve for $\boldsymbol{x}_T$. In other words, this differential equation describes a transformation from $\boldsymbol{x}_0$ to $\boldsymbol{x}_T$. In particular, this transformation is invertible: we can also solve the differential equation backward to obtain a transformation from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$. So differential equations are, in themselves, a theoretically elegant way of constructing invertible transformations.

The Jacobian Determinant

As with the diffusion models discussed previously, in this article we treat $\boldsymbol{x}_0$ as a data sample and $\boldsymbol{x}_T$ as a sample from a simple distribution. We want to use a differential equation to realize the transformation from the data distribution to the simple distribution.

First, let's understand the differential equation $\eqref{eq:ode}$ from a discretized point of view:

\begin{equation}\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t = \boldsymbol{f}_t(\boldsymbol{x}_t)\Delta t\label{eq:ode-diff}\end{equation}

Since this is a deterministic transformation, we have

\begin{equation}p_t(\boldsymbol{x}_t) d\boldsymbol{x}_t = p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) d\boldsymbol{x}_{t+\Delta t} = p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) \left| \frac{\partial \boldsymbol{x}_{t+\Delta t}}{\partial \boldsymbol{x}_t} \right| d\boldsymbol{x}_t\end{equation}

Here $\frac{\partial \boldsymbol{x}_{t+\Delta t}}{\partial \boldsymbol{x}_t}$ denotes the Jacobian matrix of the transformation, and $|\cdot|$ denotes the absolute value of its determinant. Taking the partial derivative of both sides of equation $\eqref{eq:ode-diff}$ directly, we get

\begin{equation}\frac{\partial \boldsymbol{x}_{t+\Delta t}}{\partial \boldsymbol{x}_t} = \boldsymbol{I} + \frac{\partial \boldsymbol{f}_t(\boldsymbol{x}_t)}{\partial \boldsymbol{x}_t}\Delta t\end{equation}

According to The Derivative of a Determinant, we then have

\begin{equation}\left|\frac{\partial \boldsymbol{x}_{t+\Delta t}}{\partial \boldsymbol{x}_t}\right| \approx 1 + \text{Tr}\,\frac{\partial \boldsymbol{f}_t(\boldsymbol{x}_t)}{\partial \boldsymbol{x}_t}\Delta t = 1 + \nabla_{\boldsymbol{x}_t}\cdot \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t\approx e^{\nabla_{\boldsymbol{x}_t}\cdot \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t}\end{equation}

and so we can write

\begin{equation}\log p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) - \log p_t(\boldsymbol{x}_t) \approx -\nabla_{\boldsymbol{x}_t}\cdot \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t\label{eq:approx-ode}\end{equation}

The Taylor Approximation

Suppose $p_t(\boldsymbol{x}_t)$ is the probability density function of a family of distributions that vary continuously with the parameter $t$, where $p_0(\boldsymbol{x}_0)$ is the data distribution and $p_T(\boldsymbol{x}_T)$ is the simple distribution. When $\Delta t$ and $\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t$ are both small, we have the first-order Taylor approximation

\begin{equation}\log p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) - \log p_t(\boldsymbol{x}_t) \approx (\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t)\cdot \nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t) + \Delta t\frac{\partial}{\partial t}\log p_t(\boldsymbol{x}_t)\end{equation}

Substituting $\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t$ from equation $\eqref{eq:ode-diff}$ and comparing with equation $\eqref{eq:approx-ode}$, we obtain the equation that $\boldsymbol{f}_t(\boldsymbol{x}_t)$ must satisfy:

\begin{equation}-\nabla_{\boldsymbol{x}_t}\cdot \boldsymbol{f}_t(\boldsymbol{x}_t) = \boldsymbol{f}_t(\boldsymbol{x}_t)\cdot \nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t) + \frac{\partial}{\partial t}\log p_t(\boldsymbol{x}_t)\label{eq:ode-f-eq}\end{equation}

In other words, any $\boldsymbol{f}_t(\boldsymbol{x}_t)$ satisfying this equation can be used to construct a differential equation $\eqref{eq:ode}$, and solving it realizes the transformation between the data distribution and the simple distribution. We can also rearrange it as

\begin{equation}\frac{\partial}{\partial t} p_t(\boldsymbol{x}_t) = - \nabla_{\boldsymbol{x}_t}\cdot\Big(\boldsymbol{f}_t(\boldsymbol{x}_t) p_t(\boldsymbol{x}_t)\Big)\label{eq:ode-f-eq-fp}\end{equation}

which is precisely the special case of the "Fokker-Planck equation" introduced in Rambling on About Generative Diffusion Models (6): The General Framework in the ODE Perspective, taken at $g_t=0$.

The Heat Equation

Let us consider a solution of the following form:

\begin{equation}\boldsymbol{f}_t(\boldsymbol{x}_t) = - \boldsymbol{D}_t(\boldsymbol{x}_t)\,\nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t)\label{eq:ode-f-grad}\end{equation}

where $\boldsymbol{D}_t(\boldsymbol{x}_t)$ may be a matrix or a scalar, depending on the level of complexity we're willing to consider. Why consider a solution of this form? Honestly, I initially tried to fit it into the DDIM-style formula, and only later discovered that, once generalized, it connects to the diffusion equation below—so I just directly posited it as equation $\eqref{eq:ode-f-grad}$. Looking back, if we assume $\boldsymbol{D}_t(\boldsymbol{x}_t)$ is a non-negative scalar function, then substituting it into equation $\eqref{eq:ode-diff}$ reveals a form somewhat reminiscent of gradient descent: going from $\boldsymbol{x}_0$ to $\boldsymbol{x}_T$ gradually seeks out low-probability regions, while conversely going from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$ gradually seeks out high-probability regions—which matches intuition, and can serve as a heuristic motivation for equation $\eqref{eq:ode-f-grad}$.

Substituting equation $\eqref{eq:ode-f-grad}$ into equation $\eqref{eq:ode-f-eq-fp}$, we obtain

\begin{equation}\frac{\partial}{\partial t}p_t(\boldsymbol{x}_t) = \nabla_{\boldsymbol{x}_t}\cdot\Big(\boldsymbol{D}_t(\boldsymbol{x}_t)\,\nabla_{\boldsymbol{x}_t} p_t(\boldsymbol{x}_t)\Big)\end{equation}

This is exactly the "diffusion equation" from the theory of partial differential equations. Here we only consider a very simple case: $\boldsymbol{D}_t(\boldsymbol{x}_t)$ is a scalar function $D_t$ independent of $\boldsymbol{x}_t$, in which case the diffusion equation simplifies to

\begin{equation}\frac{\partial}{\partial t}p_t(\boldsymbol{x}_t) = D_t \nabla_{\boldsymbol{x}_t}^2 p_t(\boldsymbol{x}_t)\label{eq:heat}\end{equation}

This is the "heat equation," which is the object we will focus on solving and analyzing next.

Solving for the Distribution

Using the Fourier transform, we can convert the heat equation into an ordinary differential equation, and then complete the solution for the distribution $p_t(\boldsymbol{x}_t)$, with the result:

\begin{equation}\begin{aligned} p_t(\boldsymbol{x}_t) =&\, \int \frac{1}{(2\pi\sigma_t^2)^{d/2}}\exp\left(-\frac{\Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert^2}{2\sigma_t^2}\right)p_0(\boldsymbol{x}_0) d \boldsymbol{x}_0 \\ =&\, \int \mathcal{N}(\boldsymbol{x}_t; \boldsymbol{x}_0, \sigma_t^2 \boldsymbol{I})\, p_0(\boldsymbol{x}_0) d \boldsymbol{x}_0 \end{aligned}\label{eq:heat-sol}\end{equation}

where $\sigma_t^2 = 2\int_0^t D_s ds$, or equivalently $D_t = \dot{\sigma}_t \sigma_t$ (with $\sigma_0=0$). We can see that the solution of the heat equation is precisely a Gaussian mixture model with $p_0(\boldsymbol{x}_0)$ as the initial distribution.

Derivation: Here we briefly walk through the idea behind solving the heat equation. Readers who don't care about the derivation, or who are already familiar with the heat equation, can skip this part.
Solving the heat equation $\eqref{eq:heat}$ via the Fourier transform is actually simple. Taking the Fourier transform of both sides with respect to the variable $\boldsymbol{x}_t$, and applying the rule that $\nabla_{\boldsymbol{x}_t}\to i\boldsymbol{\omega}$, the result is
\begin{equation}\frac{\partial}{\partial t}\mathcal{F}_t(\boldsymbol{\omega}) = -D_t \boldsymbol{\omega}^2 \mathcal{F}_t(\boldsymbol{\boldsymbol{\omega}})\end{equation}
This is just an ordinary differential equation in $t$, which can be solved to give
\begin{equation}\mathcal{F}_t(\boldsymbol{\omega}) = \mathcal{F}_0(\boldsymbol{\omega}) \exp\left(-\frac{1}{2}\sigma_t^2 \boldsymbol{\omega}^2\right)\end{equation}
where $\sigma_t^2 = 2\int_0^t D_s ds$, and $\mathcal{F}_0(\boldsymbol{\omega})$ is the Fourier transform of $p_0(\boldsymbol{x}_0)$. Now taking the inverse Fourier transform of both sides, $\mathcal{F}_t(\boldsymbol{\omega})$ naturally reverts to $p_t(\boldsymbol{x}_t)$, $\mathcal{F}_0(\boldsymbol{\omega})$ reverts to $p_0(\boldsymbol{x}_0)$, and $\exp\left(-\frac{1}{2}\sigma_t^2 \boldsymbol{\omega}^2\right)$ corresponds to the normal distribution $\mathcal{N}(\boldsymbol{x}_t; \boldsymbol{0}, \sigma_t^2 \boldsymbol{I})$. Finally, using the convolution property of the Fourier transform, we arrive at the solution $\eqref{eq:heat-sol}$.

Completing the Design

Let's now pull together our results. By solving the heat equation, we've determined that

\begin{equation}p_t(\boldsymbol{x}_t) = \int \mathcal{N}(\boldsymbol{x}_t; \boldsymbol{x}_0, \sigma_t^2 \boldsymbol{I})\, p_0(\boldsymbol{x}_0) d \boldsymbol{x}_0 \label{eq:heat-sol-2}\end{equation}

The corresponding differential equation at this point,

\begin{equation}\frac{d\boldsymbol{x}_t}{dt}=-\dot{\sigma}_t \sigma_t \nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t)\end{equation}

gives a deterministic transformation from $p_0(\boldsymbol{x}_0)$ to $p_T(\boldsymbol{x}_T)$. If $p_T(\boldsymbol{x}_T)$ is easy to sample from, and $\nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t)$ is known, then we can randomly sample $\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T)$ and solve the differential equation backward to generate a sample of $\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)$.

The first question is: under what conditions is $p_T(\boldsymbol{x}_T)$ easy to sample from? From the result $\eqref{eq:heat-sol-2}$, we know that

\begin{equation}\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T) \quad\Leftrightarrow\quad \boldsymbol{x}_T = \boldsymbol{x}_0 + \sigma_T \boldsymbol{\varepsilon},\,\,\, \boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\,\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\end{equation}

When $\sigma_T$ is sufficiently large, the influence of $\boldsymbol{x}_0$ on $\boldsymbol{x}_T$ becomes very weak, and we can then treat it as

\begin{equation}\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T) \quad\Leftrightarrow\quad \boldsymbol{x}_T = \sigma_T \boldsymbol{\varepsilon},\,\,\,\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\end{equation}

This achieves the goal of making $p_T(\boldsymbol{x}_T)$ easy to sample from. Hence, the general requirement for choosing $\sigma_t$ is that it be a smooth, monotonically increasing function satisfying $\sigma_0 = 0$ and $\sigma_T \gg 1$.

The second question is how to compute $\nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t)$. This is actually the same as the "score matching" section in Rambling on About Generative Diffusion Models (5): The General Framework in the SDE Perspective—we use a neural network $\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$ to fit it, with the training objective

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{x}_t \sim \mathcal{N}(\boldsymbol{x}_t; \boldsymbol{x}_0, \sigma_t^2 \boldsymbol{I})p_0(\boldsymbol{x}_0)}\left[\left\Vert \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \nabla_{\boldsymbol{x}_t} \log \mathcal{N}(\boldsymbol{x}_t; \boldsymbol{x}_0, \sigma_t^2 \boldsymbol{I})\right\Vert^2\right] \end{equation}

This is called "conditional score matching," and since we already gave its derivation in the SDE post, we won't repeat it here.

Summary

In this post, we gave a "top-down" derivation of ODE-based diffusion models: starting from the ODE, and combining it with the Jacobian determinant, we obtained a first-order approximation of the change in probability; then, comparing this with the first-order Taylor expansion directly, we derived the equation that the ODE must satisfy, which we then converted into a diffusion equation, and finally into a heat equation to solve. Relatively speaking, the whole process gets there more directly, without needing to pass through intermediate results like the SDE or the Fokker-Planck equation.