Diffusion Models Revisited (14): General Steps for Constructing an ODE (Part 1)
Continuing from where we left off: in Diffusion Models Revisited (13): From Universal Gravitation to Diffusion Models, we introduced an ODE-based generative diffusion model inspired by universal gravitation, with a very clear geometric meaning. Some readers, after reading it, raised the question: "universal gravitation" doesn't seem to be the only option — could other forms of force be used to build diffusion models under the same physical picture? On another note, while the model is indeed intuitive from a physics standpoint, it still lacks a mathematical proof that it can actually learn the data distribution in the end.
This article attempts to answer, in a fairly precise mathematical way, the question of "what kind of force field is suitable for constructing an ODE-based generative diffusion model."
Basic Result
To answer this question, we need a result we derived in Diffusion Models Revisited (12): Tackling the Diffusion ODE Head-On, concerning how the distribution changes under an ordinary differential equation.
Consider the first-order (constant) differential equation (system) for $\boldsymbol{x}_t\in\mathbb{R}^d, t\in[0,T]$
\begin{equation}\frac{d\boldsymbol{x}_t}{dt}=\boldsymbol{f}_t(\boldsymbol{x}_t)\label{eq:ode}\end{equation}
which describes an (invertible) transformation from $\boldsymbol{x}_0$ to $\boldsymbol{x}_T$. If $\boldsymbol{x}_0$ is a random variable, then $\boldsymbol{x}_t$ throughout the whole process is also a random variable, and its distributional change follows the equation
\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}
This result can be derived following the "Jacobian determinant + Taylor approximation" style used in Diffusion Models Revisited (12): Tackling the Diffusion ODE Head-On, or, as in Diffusion Models Revisited (6): The ODE Version of the General Framework, by first deriving the full "Fokker–Planck equation" and then taking $g_t=0$. Incidentally, equation $\eqref{eq:ode-f-eq-fp}$ is very well known in physics — it's called the "continuity equation" and is one manifestation of various conservation laws.
Back to diffusion models: what a diffusion model wants to do is construct a transformation that turns samples from a simple distribution into samples from the target distribution. Using equation $\eqref{eq:ode-f-eq-fp}$, in principle we can, from a given $p_t(\boldsymbol{x}_t)$, work out a feasible $\boldsymbol{f}_t(\boldsymbol{x}_t)$, and then use equation $\eqref{eq:ode}$ to complete the generation process. Note that equation $\eqref{eq:ode-f-eq-fp}$ is just a single equation, but the $\boldsymbol{f}_t(\boldsymbol{x}_t)$ we're solving for has $d$ components — so this is an underdetermined equation. In principle, we can specify essentially the entire $p_t(\boldsymbol{x}_t)$ arbitrarily (not merely the two boundary conditions $t=0,T$) in order to solve for $\boldsymbol{f}_t(\boldsymbol{x}_t)$.
So, in theory, building an ODE-based diffusion model amounts to solving an almost unconstrained, quite easy underdetermined equation. That's indeed true, but the problem is that the solutions obtained this way tend to run into practical difficulties — in plain terms, they're hard to implement in code. So the precise formulation of the problem is: how do we extract a more practically usable solution from equation $\eqref{eq:ode-f-eq-fp}$?
Simplifying the Equation
Notice that equation $\eqref{eq:ode-f-eq-fp}$ can be rewritten as
\begin{equation}\underbrace{\left(\frac{\partial}{\partial t}, \nabla_{\boldsymbol{x}_t}\right)}_{\nabla_{(t,\, \boldsymbol{x}_t)}}\cdot \underbrace{\Big(p_t( \boldsymbol{x}_t), \boldsymbol{f}_t(\boldsymbol{x}_t) p_t(\boldsymbol{x}_t)\Big)}_{\boldsymbol{u}\in\mathbb{R}^{d+1}}=0\end{equation}
As shown above, we can treat $\left(\frac{\partial}{\partial t},\nabla_{\boldsymbol{x}_t}\right)$ as exactly the $d+1$-dimensional gradient $\nabla_{(t,\, \boldsymbol{x}_t)}$, and $\big(p_t( \boldsymbol{x}_t), \boldsymbol{f}_t(\boldsymbol{x}_t) p_t(\boldsymbol{x}_t)\big)$ can be combined into a single $d+1$-dimensional vector $\boldsymbol{u}(t, \boldsymbol{x}_t)$, so $\eqref{eq:ode-f-eq-fp}$ can be written as the simple divergence equation
\begin{equation}\nabla_{(t,\, \boldsymbol{x}_t)}\cdot\boldsymbol{u}(t, \boldsymbol{x}_t)=0\label{eq:div-eq}\end{equation}
Under this form we have
\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{f}_t(\boldsymbol{x}_t) = \frac{\boldsymbol{u}_{> 1}(t, \boldsymbol{x}_t)}{\boldsymbol{u}_1(t, \boldsymbol{x}_t)}\label{eq:div-eq-ode}\end{equation}
where $\boldsymbol{u}_1$ and $\boldsymbol{u}_{> 1}$ denote respectively the first component and the remaining $d$ components of $\boldsymbol{u}$. Of course, we mustn't forget the constraint
\begin{equation}\left\{\begin{aligned} &\boldsymbol{u}_1(0, \boldsymbol{x}_0) = p_0(\boldsymbol{x}_0)\quad&(\text{initial condition}) \\[5pt] &\int \boldsymbol{u}_1(t, \boldsymbol{x}_t) d\boldsymbol{x}_t = 1\quad&(\text{integral condition}) \end{aligned}\right.\end{equation}
where $p_0(\boldsymbol{x}_0)$ is the data distribution, i.e., the target sample distribution we want to generate. As for the terminal distribution at $t=T$, our only requirement is that it be as simple as possible, so that it's easy to sample from; there's no other quantitative requirement, so we don't need to write it out explicitly here.
Green's Function
After this reformulation, we can regard $\boldsymbol{u}(t, \boldsymbol{x}_t)$ as a $d+1$-dimensional vector field, and the differential equation $\eqref{eq:div-eq-ode}$ describes exactly the trajectory of a particle moving along the field lines — which brings us right back to the physical picture presented in Diffusion Models Revisited (13): From Universal Gravitation to Diffusion Models.
To find the general solution for $\boldsymbol{u}(t, \boldsymbol{x}_t)$, we can use the idea of a Green's function. First, let's try to solve the following problem:
\begin{equation}\left\{\begin{aligned} &\nabla_{(t,\, \boldsymbol{x}_t)}\cdot\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)=0\\ &\boldsymbol{G}_1(0, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) = \delta(\boldsymbol{x}_t - \boldsymbol{x}_0),\int \boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) d\boldsymbol{x}_t = 1 \end{aligned}\right.\label{eq:div-green}\end{equation}
It's easy to show that, if the above holds, then
\begin{equation}\boldsymbol{u}(t, \boldsymbol{x}_t) = \int \boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)p_0(\boldsymbol{x}_0) d\boldsymbol{x}_0 = \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)}[\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)]\label{eq:div-green-int}\end{equation}
is a solution to equation $\eqref{eq:div-eq}$ satisfying the corresponding constraint. In this way, we've expressed $\boldsymbol{u}(t, \boldsymbol{x}_t)$ as an expectation over training samples, which is convenient for training the model. It's not hard to see that $\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ here is essentially the conditional probability $p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)$ found in diffusion models.
In fact, the $\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ defined by equation $\eqref{eq:div-green}$ isn't a Green's function in the usual sense. A typical Green's function refers to the solution under a point source, whereas here the "point source" of our Green's function is placed at the boundary. Nevertheless, the $\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ defined this way still has properties analogous to a standard Green's function — it itself is equivalent to the "force field" generated by a point source, and equation $\eqref{eq:div-green-int}$ is precisely the integral over point-source fields that yields the field of a continuous distributed source.
Universal Gravitation
Now let's use this framework to work out some concrete results. As mentioned earlier, equation $\eqref{eq:div-eq}$ or $\eqref{eq:div-green}$ is an underdetermined equation with "$d+1$ unknowns, one equation," which in theory has infinitely many solutions of all sorts. To solve it, we actually need to introduce some extra assumptions, so as to pin the solution down more precisely. The first solution is based on an isotropy assumption, and it corresponds exactly to the result in Diffusion Models Revisited (13): From Universal Gravitation to Diffusion Models.
Solving Under the Assumption
Note that "isotropy" here refers to isotropy within the $d+1$-dimensional space formed by $(t,\boldsymbol{x}_t)$, meaning that $\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ points toward the source point $(0,\,\boldsymbol{x}_0)$, and its magnitude depends only on $R = \sqrt{(t-0)^2 + \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert^2}$. Hence we can set
\begin{equation}\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) = \varphi(R)(t, \boldsymbol{x}_t - \boldsymbol{x}_0)\end{equation}
giving us
\begin{equation}\begin{aligned} 0 =&\, \nabla_{(t,\, \boldsymbol{x}_t)}\cdot\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) \\ =&\, \nabla_{(t,\, \boldsymbol{x}_t)}\varphi(R)\cdot(t, \boldsymbol{x}_t - \boldsymbol{x}_0) + \varphi(R)\nabla_{(t,\, \boldsymbol{x}_t)}\cdot (t, \boldsymbol{x}_t - \boldsymbol{x}_0) \\ =&\, \varphi'(R) \frac{(t, \boldsymbol{x}_t - \boldsymbol{x}_0)}{R}\cdot(t, \boldsymbol{x}_t - \boldsymbol{x}_0) + (d+1)\varphi(R)\\ =&\, \varphi'(R) R + (d+1)\varphi(R) \\ =&\,\frac{[\varphi(R)R^{d+1}]'}{R^d} \end{aligned}\end{equation}
that is, $[\varphi(R)R^{d+1}]'=0$, or $\varphi(R)R^{d+1}=C$, i.e., $\varphi(R)=C\times R^{-(d+1)}$. So one candidate solution is
\begin{equation}\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) = C\times\frac{(t, \boldsymbol{x}_t - \boldsymbol{x}_0)}{\left(t^2 + \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert^2\right)^{(d+1)/2}}\end{equation}
Constraints
We can see that, under the isotropy assumption, the universal-gravitation solution is unique. To prove this is a valid solution, we still need to check the constraints, the key one being
\begin{equation}\int\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) d\boldsymbol{x}_t = C\times \int\frac{t}{\left(t^2 + \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert^2\right)^{(d+1)/2}}d\boldsymbol{x}_t\end{equation}
We really just need to check that the integral result doesn't depend on $t$ or $\boldsymbol{x}_0$; then we can choose an appropriate constant $C$ to make the integral equal to 1. As for $t > 0$, we can check this by making the substitution $\boldsymbol{z} = (\boldsymbol{x}_t - \boldsymbol{x}_0) / t$: since the range of $\boldsymbol{x}_t$ covers all of space, so does the range of $\boldsymbol{z}$. Substituting this into the above gives
\begin{equation}\int\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) d\boldsymbol{x}_t = C\times \int\frac{1}{\left(1 + \Vert \boldsymbol{z}\Vert^2\right)^{(d+1)/2}}d\boldsymbol{z}\label{eq:pz}\end{equation}
We can now see the integral result no longer depends on $t$ or $\boldsymbol{x}_0$. So, as long as we choose an appropriate $C$, this check that the integral equals 1 passes. Below we assume that a $C$ making the integral 1 has already been chosen.
As for the initial value, we need to verify $\lim\limits_{t\to 0^+}\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) = \delta(\boldsymbol{x}_t - \boldsymbol{x}_0)$, which we can do simply by checking against the definition of the Dirac delta function:
1. When $\boldsymbol{x}_t\neq \boldsymbol{x}_0$, the limit is obviously 0;
2. When $\boldsymbol{x}_t = \boldsymbol{x}_0$, the limit is obviously $\infty$;
3. We've already checked that the integral of $\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ over $\boldsymbol{x}_t$ is always 1.
These three points are precisely the basic properties of the Dirac delta function — arguably even its definition — so this initial-value check also passes.
Analyzing the Result
Now, from equation $\eqref{eq:div-green-int}$ we have
\begin{equation}\boldsymbol{u}(t, \boldsymbol{x}_t) = C\times\mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)}\left[\frac{(t, \boldsymbol{x}_t - \boldsymbol{x}_0)}{\left(t^2 + \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert^2\right)^{(d+1)/2}}\right]\end{equation}
Next, we can use $\mathbb{E}_{\boldsymbol{x}}[\boldsymbol{x}] = \mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{x} - \boldsymbol{\mu}\Vert^2\right]$ to construct a score-matching-like learning objective, following a procedure we've already discussed many times, so we won't repeat it here.
As mentioned before, $\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ is actually $p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)$, and now we know its explicit form is
\begin{equation}p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)\propto \frac{t}{\left(t^2 + \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert^2\right)^{(d+1)/2}}\end{equation}
When $t=T$ is sufficiently large, the influence of $\boldsymbol{x}_0$ becomes negligible, i.e., $p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)$ degenerates into a prior distribution independent of $\boldsymbol{x}_0$:
\begin{equation}p_{prior}(\boldsymbol{x}_T) \propto \frac{T}{(T^2 + \Vert\boldsymbol{x}_T\Vert^2)^{(d+1)/2}}\end{equation}
We previously went through quite a bit of trouble to derive this result in Diffusion Models Revisited (13): From Universal Gravitation to Diffusion Models, but under this framework the result falls right into place. What's more, now that we also have $p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)$, in theory we can complete the sampling of $\boldsymbol{x}_t\sim p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)$. From the derivation of equation $\eqref{eq:pz}$, we know that, making the substitution $\boldsymbol{z} = (\boldsymbol{x}_t - \boldsymbol{x}_0) / t$, we get
\begin{equation}p(\boldsymbol{z}) \propto \frac{1}{\left(1 + \Vert \boldsymbol{z}\Vert^2\right)^{(d+1)/2}}\label{eq:pz-2}\end{equation}
So we can first sample from $p(\boldsymbol{z})$, and then obtain the corresponding $\boldsymbol{x}_t$ via $\boldsymbol{x}_t = \boldsymbol{x}_0 + t\, \boldsymbol{z}$. As for sampling from $p(\boldsymbol{z})$, since it only depends on the magnitude, we can sample the magnitude via the inverse CDF method and then randomly sample a direction to form the full sample — completely analogous to sampling from the prior distribution. However, while further investigating a leftover question below, the author stumbled upon a rather unexpected surprise!
Revisiting an Open Question
In Diffusion Models Revisited (13): From Universal Gravitation to Diffusion Models, we pointed out that the sampling scheme given in the original paper was:
\begin{equation}\boldsymbol{x}_t = \boldsymbol{x}_0 + \Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert (1+\tau)^m \boldsymbol{u},\quad t = |\varepsilon_t| (1+\tau)^m\end{equation}
where $(\boldsymbol{\varepsilon}_{\boldsymbol{x}},\varepsilon_t)\sim\mathcal{N}(\boldsymbol{0}, \sigma^2\boldsymbol{I}_{(d+1)\times(d+1)})$, $m\sim U[0,M]$, $\boldsymbol{u}$ is a unit vector uniformly distributed on the $d$-dimensional unit sphere, and $\tau,\sigma,M$ are all constants. At the time, our assessment of this sampling scheme was that it was "rather subjective" — meaning we felt it was designed subjectively by the original authors, without much justification. However, whether intentionally or not on the author's part, I discovered a remarkable "coincidence": this sampling scheme is exactly an implementation of equation $\eqref{eq:pz-2}$!
Let's prove this now. First, substitute the latter half of the above equation into the first half, giving
\begin{equation}\boldsymbol{x}_t = \boldsymbol{x}_0 + t\times \frac{\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert}{|\varepsilon_t|} \boldsymbol{u}\end{equation}
This already has the same form as $\boldsymbol{x}_t = \boldsymbol{x}_0 + t\, \boldsymbol{z}$ from the previous section, and $\boldsymbol{u}$ is also an isotropic unit random vector, so the question becomes whether $\frac{\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert}{|\varepsilon_t|}$ has the same distribution as $\Vert\boldsymbol{z}\Vert$ — and the answer is yes! Note that when converting probability density from Cartesian coordinates to spherical coordinates, we need to multiply by an extra factor of $\text{radius}^{d-1}$, so from equation $\eqref{eq:pz-2}$ we have
\begin{equation}p(\Vert\boldsymbol{z}\Vert) \propto \frac{\Vert \boldsymbol{z}\Vert^{d-1}}{\left(1 + \Vert \boldsymbol{z}\Vert^2\right)^{(d+1)/2}}\label{eq:pz-3}\end{equation}
And from $(\boldsymbol{\varepsilon}_{\boldsymbol{x}},\varepsilon_t)\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I}_{(d+1)\times(d+1)})$ (since we're studying a ratio, the variance cancels out, so for simplicity we set $\sigma=1$), we have
\begin{equation}p(\Vert\boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert) \propto \Vert\boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert^{d-1} e^{-\Vert\boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert^2/2}, \quad p(|\varepsilon_t|) \propto e^{-|\varepsilon_t|^2/2}\end{equation}
Let $r = \frac{\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert}{|\varepsilon_t|}$, then $\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert=r|\varepsilon_t|$, and then, by equating probabilities, we get
\begin{equation}\begin{aligned} p(r)dr =&\, \mathbb{E}_{|\varepsilon_t|\sim p(|\varepsilon_t|)}\big[p(\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert\color{red}{=r|\varepsilon_t|})d(\color{red}{r|\varepsilon_t|})\big] \\[5pt] \propto&\, \mathbb{E}_{|\varepsilon_t|\sim p(|\varepsilon_t|)}\big[r^{d-1}|\varepsilon_t|^d e^{-r^2|\varepsilon_t|^2/2} dr\big] \\[5pt] \propto&\, \int_0^{\infty} r^{d-1}|\varepsilon_t|^d e^{-r^2|\varepsilon_t|^2/2} e^{-|\varepsilon_t|^2/2} d|\varepsilon_t| dr \\ =&\, \int_0^{\infty} r^{d-1}|\varepsilon_t|^d e^{-(r^2+1)|\varepsilon_t|^2/2} d|\varepsilon_t| dr \\ =&\, \frac{r^{d-1}}{(1+r^2)^{(d+1)/2}} \int_0^{\infty} s^d e^{-s^2/2} ds dr \quad\left(\text{let}s = |\varepsilon_t|\sqrt{r^2+1}\right) \\ \propto&\, \frac{r^{d-1}}{(1+r^2)^{(d+1)/2}} dr \end{aligned}\end{equation}
Hence $p(r)\propto \frac{r^{d-1}}{(1+r^2)^{(d+1)/2}}$, which is exactly consistent with $\eqref{eq:pz-3}$. So $\frac{\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert}{|\varepsilon_t|}\boldsymbol{u}$ does indeed provide a valid way to sample $\boldsymbol{z}$ — one that's much simpler to implement than the inverse CDF method, though the original paper never mentioned this.
Separating Space and Time
We've just solved the isotropic solution in the $d+1$-dimensional space formed by $(t,\boldsymbol{x}_t)$, which is, in a sense, the simplest possible solution. Some readers might find this hard to accept, given that this universal-gravitation diffusion model looks noticeably more complex mathematically. But in fact, when solving mathematical-physics equations, isotropic solutions are very often tried first precisely because they're the simplest ones.
Of course, treating $(t,\boldsymbol{x}_t)$ as isotropic over the combined "space-time," while mathematically valid, isn't the most intuitive way to think about it. We're more used to thinking of isotropy in space alone, treating the time dimension separately. This section solves the problem under that assumption.
Solving Under the Assumption
That is, in this part "isotropy" refers to isotropy in the $d$-dimensional space of $\boldsymbol{x}_t$, where $\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ is decomposed into two parts, $(\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0), \boldsymbol{G}_{> 1}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0))$, for the purposes of interpretation. Here $\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ is just a scalar; isotropy means it only depends on $r = \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert$, which we denote $\phi_t(r)$. And $\boldsymbol{G}_{> 1}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ is a $d$-dimensional vector; isotropy means $\boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$ points toward the source point $\boldsymbol{x}_0$, and its magnitude depends only on $r = \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert$. Hence we can set
\begin{equation}\boldsymbol{G}_{>1}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) = \varphi_t(r)(\boldsymbol{x}_t - \boldsymbol{x}_0)\end{equation}
giving us
\begin{equation}\begin{aligned} 0 =&\, \frac{\partial}{\partial t}\phi_t(r) + \nabla_{\boldsymbol{x}_t}\cdot(\varphi_t(r) (\boldsymbol{x}_t - \boldsymbol{x}_0)) \\ =&\, \frac{\partial}{\partial t}\phi_t(r) + r\frac{\partial}{\partial r}\varphi_t(r) + d\, \varphi_t(r) \\ =&\, \frac{\partial}{\partial t}\phi_t(r) + \frac{1}{r^{d-1}}\frac{\partial}{\partial r}\big(\varphi_t(r) r^d\big)\\ \end{aligned}\end{equation}
Here we have two functions to determine, $\phi_t(r)$ and $\varphi_t(r)$, but only one equation, so solving is even easier. Since the constraint constrains $\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0)$, i.e., $\phi_t(r)$ rather than $\varphi_t(r)$, for simplicity we usually pick a $\phi_t(r)$ satisfying the condition and then solve for $\varphi_t(r)$, giving the result
\begin{equation}\varphi_t(r) = -\frac{1}{r^d}\int \frac{\partial}{\partial t}\phi_t(r) r^{d-1} dr = -\frac{1}{r^d}\frac{\partial}{\partial t}\int \phi_t(r) r^{d-1} dr\label{eq:f-g-t-r}\end{equation}
Gaussian Diffusion
In this part, let's show that the common ODE diffusion models based on the Gaussian assumption are also a special case of equation $\eqref{eq:f-g-t-r}$. For the Gaussian assumption, we have
\begin{equation}\boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) = p_t(\boldsymbol{x}_t|\boldsymbol{x}_0) = \frac{1}{(2\pi\sigma_t^2)^{d/2}} e^{-\Vert\boldsymbol{x}_t-\boldsymbol{x}_0\Vert^2/2\sigma_t^2}\end{equation}
that is, $\phi_t(r) = \frac{1}{(2\pi\sigma_t^2)^{d/2}} e^{-r^2/2\sigma_t^2}$, where $\sigma_t$ is a function of $t$ that increases monotonically, satisfying $\sigma_0=0$, and for sufficiently large $\sigma_T$, $\sigma_0=0$, so that the initial condition holds; sufficiently large $\sigma_T$ ensures the prior distribution is independent of the data. As for the constraint that the integral equal 1, that's a basic property of the Gaussian distribution and is naturally satisfied.
Substituting into equation $\eqref{eq:f-g-t-r}$ and solving gives:
\begin{equation}\varphi_t(r) = \frac{\dot{\sigma}_t}{(2\pi\sigma_t^2)^{d/2}\sigma_t} e^{-r^2/2\sigma_t^2} = \frac{\dot{\sigma}_t}{\sigma_t}\phi_t(r)\end{equation}
where the integral for $r$ involves the incomplete gamma function and is fairly complicated — the author computed it directly with Mathematica. With this result in hand, we have
\begin{equation}\begin{aligned} \boldsymbol{u}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) =&\, \int p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)p_0(\boldsymbol{x}_0) d\boldsymbol{x}_0 = p_t(\boldsymbol{x}_t) \\ \boldsymbol{u}_{> 1}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) =&\, \int \frac{\dot{\sigma}_t}{\sigma_t}(\boldsymbol{x}_t - \boldsymbol{x}_0)p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)p_0(\boldsymbol{x}_0) d\boldsymbol{x}_0 \\ =&\, -\dot{\sigma}_t\sigma_t \int\nabla_{\boldsymbol{x}_t} p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)p_0(\boldsymbol{x}_0) d\boldsymbol{x}_0 \\ =&\, -\dot{\sigma}_t\sigma_t \nabla_{\boldsymbol{x}_t} p_t(\boldsymbol{x}_t) \\ \end{aligned}\end{equation}
and thus, from equation $\eqref{eq:div-eq-ode}$,
\begin{equation}\boldsymbol{f}_t(\boldsymbol{x}_t) = \frac{\boldsymbol{u}_{> 1}(t, \boldsymbol{x}_t)}{\boldsymbol{u}_1(t, \boldsymbol{x}_t)} = -\dot{\sigma}_t\sigma_t \nabla_{\boldsymbol{x}_t} \log p_t(\boldsymbol{x}_t) \end{equation}
These results are exactly consistent with those in Diffusion Models Revisited (12): Tackling the Diffusion ODE Head-On; for the remaining processing details, one can refer to that article as well.
A Reverse Construction
The approach just described — given $\phi_t(r)$, solve for $\varphi_t(r)$ — is simple in theory but runs into two practical difficulties: 1) $\phi_t(r)$ needs to satisfy both the initial condition and the integral condition, which isn't always easy to construct; 2) the integral over $r$ doesn't necessarily have a simple closed form. Given this, we can think of a reverse-construction approach instead.
We know that $\phi_t(r)$ is a probability density in Cartesian coordinates, and converting to spherical coordinates requires multiplying by $C_d r^{d-1}$, where $C_d$ is some constant (related to $d$). By equation $\eqref{eq:div-eq-ode}$, the final result is a ratio and thus unaffected by constants, so for simplicity we can ignore this constant. Once we do, the integrand happens to be exactly that of equation $\eqref{eq:f-g-t-r}$, so the integral in equation $\eqref{eq:f-g-t-r}$
\begin{equation}\int \phi_t(r) r^{d-1} dr\end{equation}
is precisely a cumulative probability function (more precisely, it's $1/C_d$ of the cumulative probability function plus an unimportant constant we've already dropped). Now, computing the cumulative probability from a probability density isn't always easy, but computing the density from the cumulative probability is simple (just differentiate). So we can first construct the cumulative probability function, and then derive the corresponding $\phi_t(r),\varphi_t(r)$ from it — thereby avoiding the difficulty of integration altogether.
Specifically, we construct a cumulative probability function $\psi_t(r)$ satisfying the following conditions:
1. $\psi_t(0)=0$, $\psi_t(\infty)=1$;
2. $\psi_t(r)$ is monotonically increasing in $r$;
3. $\forall r > 0, \lim\limits_{t\to 0^+} \psi_t(r)=1$.
For anyone who's studied activation functions a bit, it shouldn't be hard to construct a function satisfying these conditions — it's essentially a smooth approximation to the "step function," such as $\tanh\left(\frac{r}{t}\right)$, $1-e^{-r/t}$, and so on. Once we have $\psi_t(r)$, from equation $\eqref{eq:f-g-t-r}$ we get
\begin{equation}\phi_t(r) = \frac{1}{r^{d-1}}\frac{\partial}{\partial r}\psi_t(r), \quad \varphi_t(r) = -\frac{1}{r^d}\frac{\partial}{\partial t}(\psi_t(r)\color{skyblue}{+\lambda_t})\end{equation}
where $\color{skyblue}{\lambda_t}$ is an arbitrary function of $t$, and can generally just be set to 0. Of course, all these isotropic solutions are essentially equivalent, including the "universal-gravitation diffusion" derived in the previous section — they can all be subsumed under the formula above, and can also be derived from one another via coordinate transformations. This is because the formula above depends only on a one-dimensional cumulative probability function $\psi_t(r)$, and the cumulative probability functions of different distributions can generally be transformed into one another (they're all well-behaved, monotonically increasing functions).
Summary
This article constructed a general framework for ODE-based diffusion. In theory, every ODE-based diffusion model can be subsumed under this framework, and from it we can derive all sorts of novel, even exotic, ODE-based diffusion models. For instance, all the derivations here are based on an isotropy assumption; in principle, one could replace the isotropic $\varphi(R)$ with a more general $\varphi(t;\boldsymbol{x}_t,\boldsymbol{x}_0)$, which could be solved using the method of The Method of Characteristics for First-Order PDEs, yielding a new family of models. All in all, this is a genuine "production workshop" for ODE-based diffusion models.
Some readers might ask: I just want a usable generative diffusion model — what's the point of all these fancy, elaborate variants? In fact, just as with A Brief Introduction to f-GAN: A Production Workshop for GAN Models and Designing GANs: Yet Another GAN Production Workshop, our hope is to discover and master the underlying rules governing the construction of generative models, so as to gain a deeper understanding of what's essential to generative models — and thereby discover more effective ones. This is an endless pursuit of perfection.
Experimental results in the earlier "universal-gravitation diffusion" paper have already shown that, as an ODE-based diffusion model, it performs somewhat better than Gaussian diffusion. This tells us that, even under the same isotropy assumption, these mathematically equivalent diffusion models can still differ in practical performance. So, how to better combine experimental details to answer the question "what kind of design makes for a better diffusion model" will be a very meaningful research direction going forward.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.