Diffusion Models Talk (17): The General Procedure for Constructing ODEs (Part III)

生成扩散模型漫谈（十七）：构建ODE的一般步骤（下） 苏剑林 · 2023-02-23 · Machine Learning probability differential equation generative model diffusion

History has an uncanny way of repeating itself. Back when I wrote Diffusion Models Talk (14): The General Procedure for Constructing ODEs (Part I) (which at the time did not have the "Part I" suffix), I thought I had already worked out the general procedure for constructing ODE-based diffusion. Then reader @gaohuazuo came up with a new, intuitive, and effective scheme, which directly led to the follow-up post Diffusion Models Talk (14): The General Procedure for Constructing ODEs (Part II) (which at the time carried the suffix "Part II" — the final part, or so I thought). Just when I believed the matter had finally been settled, I discovered that the ICLR 2023 paper Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow had come up with yet another scheme for constructing ODE-based diffusion models — one so simple and intuitive that it is truly unprecedented, the kind of thing that makes you want to applaud. So I had no choice but to quietly rename the previous post's suffix to "Part II," and write this "Part III" to share this new result.

An Intuitive Result

As we know, a diffusion model is an evolution process over $\boldsymbol{x}_T\to \boldsymbol{x}_0$, and ODE-based diffusion models specify this evolution process via the following ODE:

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

The problem of constructing an ODE-based diffusion model amounts to designing a function $\boldsymbol{f}_t(\boldsymbol{x}_t)$ such that its corresponding evolution trajectory forms a transformation between the given distributions $p_T(\boldsymbol{x}_T)$ and $p_0(\boldsymbol{x}_0)$. In plain terms, we want that if we randomly sample some $\boldsymbol{x}_T$ from $p_T(\boldsymbol{x}_T)$, then evolve it backward according to the above ODE, the resulting $\boldsymbol{x}_0$ follows $\sim p_0(\boldsymbol{x}_0)$. more

The idea in the original paper is extremely simple. We randomly pick $\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T)$, and assume they are transformed according to the trajectory

\begin{equation}\boldsymbol{x}_t = \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)\label{eq:track}\end{equation}

This trajectory is a known function — a part we design ourselves. In principle, any continuous function satisfying

\begin{equation}\boldsymbol{x}_0 = \boldsymbol{\varphi}_0(\boldsymbol{x}_0, \boldsymbol{x}_T),\quad \boldsymbol{x}_T = \boldsymbol{\varphi}_T(\boldsymbol{x}_0, \boldsymbol{x}_T)\end{equation}

will do. We can then write down the differential equation it satisfies:

\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}\label{eq:fake-ode}\end{equation}

But this differential equation is not practically usable, because what we want is to generate $\boldsymbol{x}_0$ given $\boldsymbol{x}_T$, whereas the right-hand side here is a function of $\boldsymbol{x}_0$ (if we already knew $\boldsymbol{x}_0$, we'd be done). Only an ODE like equation $\eqref{eq:ode}$, whose right-hand side involves only $\boldsymbol{x}_t$ (from a purely causal standpoint, it could in principle also include $\boldsymbol{x}_T$, but we generally don't consider that case), can be practically evolved. So here comes an intuitive yet "wildly imaginative" idea: learn a function $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$ that approximates the right-hand side above as closely as possible! To this end, we optimize the following objective:

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T)}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}\right\Vert^2\right] \label{eq:objective} \end{equation}

Since $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$ closely approximates $\frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}$, we take it that replacing the right-hand side of equation $\eqref{eq:fake-ode}$ with $\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$ is also valid, giving us the practically usable diffusion ODE:

\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\label{eq:s-ode}\end{equation}

A Simple Example

As a simple example, let us set $T=1$, and let the transformation trajectory be a straight line:

\begin{equation}\boldsymbol{x}_t = \boldsymbol{\varphi}_t(\boldsymbol{x}_0,\boldsymbol{x}_1) = (\boldsymbol{x}_1 - \boldsymbol{x}_0)t + \boldsymbol{x}_0\end{equation}

Then

\begin{equation}\frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t} = \boldsymbol{x}_1 - \boldsymbol{x}_0\end{equation}

so the training objective $\eqref{eq:objective}$ becomes:

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}\big((\boldsymbol{x}_1 - \boldsymbol{x}_0)t + \boldsymbol{x}_0, t\big) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\end{equation}

or, equivalently,

\begin{equation}\mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{x}_t\sim p_0(\boldsymbol{x}_0)p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \frac{\boldsymbol{x}_t - \boldsymbol{x}_0}{t}\right\Vert^2\right]\end{equation}

And that's it! The result is exactly the same as the "straight-line trajectory" example in Diffusion Models Talk (14): The General Procedure for Constructing ODEs (Part II), and it is also the main model studied in the original paper, called "Rectified Flow."

Looking at this straight-line example, we can see that the steps involved in constructing a diffusion ODE via this approach take only a few lines — a massive simplification compared to the previous derivations, so simple in fact that it gives one the incredible feeling of having one's whole impression of diffusion models overturned.

The Proof

So far, however, the conclusions of the "Intuitive Result" section above are only an intuitive conjecture, since we have not yet theoretically proven that the equation $\eqref{eq:s-ode}$ obtained by optimizing the objective $\eqref{eq:objective}$ indeed realizes a transformation between the distributions $p_T(\boldsymbol{x}_T)$ and $p_0(\boldsymbol{x}_0)$.

To prove this, my first instinct was to show that the optimal solution of the objective $\eqref{eq:objective}$ satisfies the continuity equation:

\begin{equation}\frac{\partial p_t(\boldsymbol{x}_t)}{\partial t} = -\nabla_{\boldsymbol{x}_t}\cdot\big(p_t(\boldsymbol{x}_t)\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\big)\end{equation}

If it does, then by the correspondence between the continuity equation and the ODE (see Diffusion Models Talk (12): "Brute-Forcing" the Diffusion ODE and Deriving the Continuity Equation and the Fokker–Planck Equation via Test Functions), equation $\eqref{eq:s-ode}$ indeed realizes a transformation between the distributions $p_T(\boldsymbol{x}_T)$ and $p_0(\boldsymbol{x}_0)$.

But on closer thought, this approach seems a bit roundabout, because according to Deriving the Continuity Equation and the Fokker–Planck Equation via Test Functions, the continuity equation itself is derived from the ODE via

\begin{equation}\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] = \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t)\Delta t)\right]\label{eq:base}\end{equation}

so, by rights, $\eqref{eq:base}$ is the more fundamental object, and we only need to prove that the optimal solution of $\eqref{eq:objective}$ satisfies it. That is, we want to find a function $\boldsymbol{f}_t(\boldsymbol{x}_t)$ that is purely a function of $\boldsymbol{x}_t$ and satisfies $\eqref{eq:base}$, and then show that it is exactly the optimal solution of $\eqref{eq:objective}$.

So we write (for brevity, $\boldsymbol{\varphi}_t(\boldsymbol{x}_0,\boldsymbol{x}_T)$ is shortened to $\boldsymbol{\varphi}_t$):

\begin{equation}\begin{aligned} \mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] =&\, \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\phi(\boldsymbol{\varphi}_{t+\Delta t})\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\phi(\boldsymbol{\varphi}_t) + \Delta t\,\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{\varphi}_t}\phi(\boldsymbol{\varphi}_t)\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\ \end{aligned}\end{equation}

where the first equality follows from equation $\eqref{eq:track}$, the second is a first-order Taylor expansion, the third is again equation $\eqref{eq:track}$, and the fourth follows because $\boldsymbol{x}_t$ is a deterministic function of $\boldsymbol{x}_0,\boldsymbol{x}_T$, so the expectation over $\boldsymbol{x}_0,\boldsymbol{x}_T$ equals the expectation over $\boldsymbol{x}_t$.

We see that $\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$ is a function of $\boldsymbol{x}_0,\boldsymbol{x}_T$. Next we make one more assumption: equation $\eqref{eq:track}$ is invertible with respect to $\boldsymbol{x}_T$. This assumption means we can solve for $\boldsymbol{x}_T=\boldsymbol{\psi}_t(\boldsymbol{x}_0,\boldsymbol{x}_t)$ from equation $\eqref{eq:track}$, and substitute this result into $\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$, turning it into a function of $\boldsymbol{x}_0,\boldsymbol{x}_t$. So we have

\begin{equation}\begin{aligned} \mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_t}\left[\underbrace{\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]}_{\boldsymbol{x}_t\text{function of}}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\ =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi\left(\boldsymbol{x}_t + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\right)\right] \end{aligned}\end{equation}

where the second equality holds because $\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$ has already been rewritten as a function of $\boldsymbol{x}_0,\boldsymbol{x}_t$, so the random variable in the second expectation term is changed to $\boldsymbol{x}_0,\boldsymbol{x}_t$; the third equality corresponds to the decomposition $p(\boldsymbol{x}_0,\boldsymbol{x}_t)=p(\boldsymbol{x}_0|\boldsymbol{x}_t)p(\boldsymbol{x}_t)$, in which $\boldsymbol{x}_0,\boldsymbol{x}_t$ is no longer independent, so we need to note $\boldsymbol{x}_0|\boldsymbol{x}_t$, i.e., that $\boldsymbol{x}_0$ depends on $\boldsymbol{x}_t$. Note that $\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$ was originally a function of $\boldsymbol{x}_0,\boldsymbol{x}_t$, but now, after taking the expectation over $\boldsymbol{x}_0$, the only remaining free variable is $\boldsymbol{x}_t$ — as we shall see, this is exactly the function that is purely a function of $\boldsymbol{x}_t$ that we've been looking for! The fourth equality simply uses the Taylor expansion formula to recombine the two terms.

We now arrive at

\begin{equation}\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] = \mathbb{E}_{\boldsymbol{x}_t}\left[\phi\left(\boldsymbol{x}_t + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\right)\right]\end{equation}

which holds for any test function $\phi$, and therefore this means

\begin{equation}\boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\quad\Rightarrow\quad\frac{d\boldsymbol{x}_t}{dt} = \mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\label{eq:real-ode}\end{equation}

is exactly the ODE we were seeking. From

\begin{equation}\mathbb{E}_{\boldsymbol{x}}[\boldsymbol{x}] = \mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{x} - \boldsymbol{\mu}\Vert^2\right]\label{eq:mean-opt}\end{equation}

the right-hand side of equation $\eqref{eq:real-ode}$ is precisely the optimal solution of the training objective $\eqref{eq:objective}$, which proves that the equation $\eqref{eq:s-ode}$ obtained by optimizing the training objective $\eqref{eq:objective}$ indeed realizes a transformation between the distributions $p_T(\boldsymbol{x}_T)$ and $p_0(\boldsymbol{x}_0)$.

Reflections

Regarding the idea for constructing the diffusion ODE described in "Intuitive Result," the authors of the original paper also wrote a Zhihu column article, [ICLR2023] A New Method for Diffusion Generative Models: Extremely Simplified, Generated in One Step, which readers are also encouraged to check out. It was in fact through this column article that I first learned of this method, and I was deeply struck and impressed by it.

If you have read Diffusion Models Talk (14): The General Procedure for Constructing ODEs (Part II), you will appreciate all the more just how simple and direct this idea is, and understand better why I am so lavish with my praise. I don't mind admitting, somewhat embarrassingly, that while writing "Part II" (then called "Part III"), I had actually considered the trajectory described by equation $\eqref{eq:track}$, but within the framework I had at the time, I simply couldn't push the derivation through, and ultimately gave up. I could never have imagined back then that it could be carried through in such a remarkably simple way. So writing this diffusion-ODE series has really given me a feeling of "comparing yourself to others will only make you furious" — "Part II" and "Part III" stand as vivid testimony to my intelligence being repeatedly "outclassed by a dimensional strike."

Readers might wonder whether there will be an even simpler fourth installment, subjecting me to yet another such humbling experience. Perhaps — but the probability really is quite small; it's genuinely hard to imagine a construction procedure simpler than this one. The "Intuitive Result" section may look long, but it really only involves two steps: (1) pick an arbitrary transition trajectory; (2) use a function of $\boldsymbol{x}_t$ to approximate the derivative of the transition trajectory with respect to $t$. With just these two steps, how could it possibly be simplified any further? Even the derivation in the "Proof" section is quite straightforward — although it's written at some length, in essence it's just taking a derivative and then changing the distribution over which the expectation is taken, which is simpler than the previous two installments by more than just a little. In short, anyone who has personally worked through the derivations in the first two installments of the ODE-diffusion series will deeply feel just how simple this approach really is — simple to the point that it feels like it truly cannot be simplified any further.

Beyond offering a simple approach to constructing the diffusion ODE, the original paper also discusses the connection between Rectified Flow and optimal transport, and how this connection can be leveraged to speed up the sampling process, among other things. This content isn't the main focus of this post, so we'll leave it for discussion another time, when the opportunity arises.

Summary

This post has introduced an extremely simple and intuitive approach to constructing ODE-based diffusion models, as proposed in the Rectified Flow paper, and has presented my own proof of it.