Generative Diffusion Models (14): General Steps for Constructing ODEs (Part 1)

Solving under Assumptions

Note that "isotropic" here refers to isotropy in the (d+1)-dimensional space composed of (t, \boldsymbol{x}_t). This means \boldsymbol{G}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) points towards 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}. Thus, 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} Then: \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, which gives \varphi(R)=C\times R^{-(d+1)}. Therefore, a 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

As we can see, under the isotropic assumption, the universal gravitation solution is the unique solution. To prove it is a feasible solution, we must 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} In fact, we only need to verify that the integral result is independent of t and \boldsymbol{x}_0; then we can choose an appropriate constant C to make the integral equal to 1. For t > 0, we can perform a change of variables \boldsymbol{z} = (\boldsymbol{x}_t - \boldsymbol{x}_0) / t. Since the range of \boldsymbol{x}_t is the entire space, \boldsymbol{z} is also over the entire space. Substituting this into the equation 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} Now it is clear that the integral result is independent of t and \boldsymbol{x}_0. Thus, by choosing an appropriate C, the condition that the integral equals 1 is satisfied. Henceforth, we assume C has been chosen such that the integral is 1.

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). This can be checked according to the definition of the Dirac delta function:

1. When \boldsymbol{x}_t \neq \boldsymbol{x}_0, the limit is clearly 0;
2. When \boldsymbol{x}_t = \boldsymbol{x}_0, the limit is clearly \infty;
3. We have already verified 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 exactly the basic properties of the Dirac delta function, or even its definition, so the initial value check also passes.

Analysis of Results

Now, according to Equation [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 objective for learning. This process has been discussed many times and will not be repeated 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). We now know its specific 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 large enough, the influence of \boldsymbol{x}_0 becomes negligible, and 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} Previously, in Part 13, deriving this result was quite involved, but within this framework, it follows naturally. Furthermore, since we now have p_t(\boldsymbol{x}_t|\boldsymbol{x}_0), we can theoretically sample \boldsymbol{x}_t \sim p_t(\boldsymbol{x}_t|\boldsymbol{x}_0). From the derivation of Equation [eq:pz], we know that if we substitute \boldsymbol{z} = (\boldsymbol{x}_t - \boldsymbol{x}_0) / t, we have: \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} Thus, 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}), it only depends on the magnitude, so we can use the inverse cumulative distribution function method to sample the magnitude and then randomly sample a direction to construct the result, which is identical to sampling from the prior distribution. However, while further investigating the following legacy problem, I discovered an unexpected "surprise"!

Problem Revisited

In Part 13, 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 constants. At the time, the evaluation of this sampling was that it had "a lot of subjectivity," meaning it felt like the author designed it subjectively without much justification. However, whether the author intended it or not, I found a magical "coincidence": this sampling is exactly an implementation of Equation [eq:pz-2]!

Let’s prove this. First, substitute the second half of the above equation into the first half: \begin{equation} \boldsymbol{x}_t = \boldsymbol{x}_0 + t\times \frac{\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert}{|\varepsilon_t|} \boldsymbol{u} \end{equation} Formally, this is already the same as \boldsymbol{x}_t = \boldsymbol{x}_0 + t\, \boldsymbol{z} mentioned in the previous section, and \boldsymbol{u} is 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. The answer is yes! Note that when the probability density changes from Cartesian coordinates to spherical coordinates, it must be multiplied by \text{radius}^{d-1}. Thus, according to Equation [eq:pz-2]: \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} Given (\boldsymbol{\varepsilon}_{\boldsymbol{x}},\varepsilon_t)\sim\mathcal{N}(\boldsymbol{0}, \boldsymbol{I}_{(d+1)\times(d+1)}) (since we are studying a ratio, the variance cancels out, so we take \sigma=1 for simplicity): \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|. Based on the equality of probabilities: \begin{equation} \begin{aligned} p(r)dr =&\, \mathbb{E}_{|\varepsilon_t|\sim p(|\varepsilon_t|)}\big[p(\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert = r|\varepsilon_t|)d(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} Therefore p(r)\propto \frac{r^{d-1}}{(1+r^2)^{(d+1)/2}}, which is perfectly consistent with [eq:pz-3]. Thus, \frac{\Vert \boldsymbol{\varepsilon}_{\boldsymbol{x}}\Vert}{|\varepsilon_t|}\boldsymbol{u} indeed provides an effective way to sample \boldsymbol{z}. This is much simpler to implement than the inverse CDF method, though the original paper did not mention this.

Solving under Assumptions

That is to say, "isotropy" here refers to isotropy in the d-dimensional space of \boldsymbol{x}_t. \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)). Here \boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) is just a scalar, and isotropy means it depends only on r = \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert, which we denote as \phi_t(r). \boldsymbol{G}_{> 1}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) is a d-dimensional vector, and isotropy means \boldsymbol{G}_{> 1}(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0) points towards the source point \boldsymbol{x}_0, and its magnitude depends only on r = \Vert \boldsymbol{x}_t - \boldsymbol{x}_0\Vert. Thus, 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} Then: \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} There are two unknown functions \phi_t(r) and \varphi_t(r), but only one equation, so solving is even simpler. Since the constraints apply to \boldsymbol{G}_1(t, 0; \boldsymbol{x}_t, \boldsymbol{x}_0), i.e., \phi_t(r) rather than \varphi_t(r), we usually specify a \phi_t(r) that satisfies the conditions and solve for \varphi_t(r). The result is: \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, we show that the common ODE diffusion models based on the Gaussian distribution assumption are also a special case of Equation [eq:f-g-t-r]. For the Gaussian assumption: \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 monotonically increasing function of t, satisfying \sigma_0=0 and \sigma_T being large enough. \sigma_0=0 is to satisfy the initial condition, and a large \sigma_T ensures the prior distribution is independent of the data. As for the constraint that the integral equals 1, this is a basic property of the Gaussian distribution and is naturally satisfied.

Substituting into Equation [eq:f-g-t-r] yields: \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} The integral over r involves the incomplete gamma function and is quite complex; I calculated it directly using Mathematica. With this result, 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} Thus, according to Equation [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 perfectly consistent with those in Part 12. For the remaining processing details, one can refer to that article.

Reverse Construction

The approach of specifying \phi_t(r) to solve for \varphi_t(r) is theoretically simple, but in practice, it faces two difficulties: 1. \phi_t(r) must satisfy both the initial condition and the integral condition, which is not easy to construct; 2. The integral over r may not have a simple elementary form. Given this, we can think of a reverse construction method.

We know that \phi_t(r) is the probability density in Cartesian coordinates. Converting to spherical coordinates requires multiplying by C_d r^{d-1}, where C_d is a constant (related to d). According to Equation [eq:div-eq-ode], the final result is a ratio and is unaffected by constants. So for simplicity, we ignore this constant. After ignoring the constant, it is exactly the integrand in Equation [eq:f-g-t-r]. Therefore, the integral in Equation [eq:f-g-t-r]: \begin{equation} \int \phi_t(r) r^{d-1} dr \end{equation} is exactly a cumulative probability function (more accurately, 1/C_d of the cumulative probability function plus a constant, but we have ignored the irrelevant constants). Calculating the cumulative probability from the density is not always easy, but calculating the density from the cumulative probability is simple (differentiation). Thus, we can first construct the cumulative probability function and then find the corresponding \phi_t(r) and \varphi_t(r), thereby avoiding the difficulty of integration.

Specifically, 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 with respect to r;
3. \forall r > 0, \lim\limits_{t\to 0^+} \psi_t(r)=1.

Students who have studied activation functions should find it easy to construct functions satisfying these conditions; they are essentially smooth approximations of the Heaviside step function, such as \tanh\left(\frac{r}{t}\right), 1-e^{-r/t}, etc. With \psi_t(r), according to Equation [eq:f-g-t-r], we have: \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) + \lambda_t) \end{equation} where \lambda_t is any function of t, which can generally be set to 0. Of course, these isotropic solutions are essentially equivalent, including the "universal gravitation diffusion" derived in the previous section. They can all be incorporated into the above formula and derived from each other through coordinate transformations. This is because the formula only depends on a univariate cumulative probability function \psi_t(r), and cumulative probability functions between different distributions can generally be transformed into each other (as they are well-behaved monotonically increasing functions).