From Denoising Diffusion to ODEs

Setting & Goals

Consider an empirical data distribution ${\pi }_{\mathrm{data}}$. In order to simulate approximate samples from ${\pi }_{\mathrm{data}}$, Denoising Diffusion Probabilistic Models (DDPM) simulate a forward diffusion process $\{X_t{\}}_{[0,T]}$ on an interval $[0,T]$. The diffusion is initialized at the data distribution, i.e. $X_0\sim {\pi }_{\mathrm{data}}$, and is chosen so that that the distribution of $X_T$ is very close to a known and tractable reference distribution ${\pi }_{\mathrm{ref}}$, e.g. a Gaussian distribution. Denote by $p_t(dx)$ the marginal distribution at time $0\le t\le T$, i.e. $\mathbb{P}(X_t\in dx)=p_t(dx)$. By choosing the forward distribution with simple and tractable transition probabilities, e.g. an Ornstein-Uhlenbeck, it is relatively easy to estimate $\mathrm{\nabla }\log p_t(x)$ from simulated data: this can be formulated as a simple regression problem. This allows one to simulate the diffusion backward in time and generate approximate samples from ${\pi }_{\mathrm{data}}$. Why this is useful is another question…

The fact that the mapping from data-samples at time $t=0$ to (approximate) Gaussian samples at time $t=T$ is stochastic and described by diffusion processes is cumbersome. This would be much more convenient to build a deterministic mapping between the data-distribution ${\pi }_{\mathrm{data}}$ and the Gaussian reference distribution ${\pi }_{\mathrm{ref}}$: this would allows one to associate a likelihood to data samples and to easily “encode”/“decode” data-samples. To so this, one can try to replace diffusion by Ordinary Differential Equations.

The diffusion-ODE trick

Consider an arbitrary diffusion process $dX=\mu (t,X)dt+dB$ with associated distribution $p_t(dx)$ at time $t$. The Fokker-Planck equation that describes the evolution of $p_t$ reads

$${\partial }_t{p}_t=-\mathrm{div}(p_t\mu )+\frac{1}{2}\mathrm{\Delta }{p}_t$$

If there were no diffusion term $dB$ and $X$ was describing instead the evolution of differential equation $dX/dt=F(t,X)$, the associated evolution of the density of $X$ would simply read

$$\begin{array}{}\text{(1)}& {\partial }_t{p}_t=-\mathrm{div}(p_{t}F).\end{array}$$

If one can find a vector field $F(t,x)$ such that

$$-\mathrm{div}(p_t\mu )+\frac{1}{2}\mathrm{\Delta }{p}_t=-\mathrm{div}(p_{t}F),$$

then one can basically replace diffusions by ODEs. The diffusion-ODE trick is the simple remark that

$$F(t,x)=\mu (t,x)-\frac{1}{2}{\mathrm{\nabla }}_x\log p_t(x)$$

does exactly this, as algebra immediately shows it. The additional term $-\frac{1}{2}{\mathrm{\nabla }}_x\log p_t(x)$ is intuitive. The coefficient $1/2$ is because one is trying to match the term $(1/2)\mathrm{\Delta }{p}_t$ in the Fokker-Planck equation. And the overall term $-\frac{1}{2}{\mathrm{\nabla }}_x\log p_t(x)$ is just driving the ODE in direction where the probability density $p_t$ is small, i.e. it follows the negative gradient of the log-density: it is exactly trying to imitate the diffusion term $dB$.

What this means is that a diffusion process $dX=\mu (t,X)dt+dB$ started from $X_0\sim p_0$ and marginal distribution $X_t\sim p_t(dx)$ can be imitated by an ODE process $dY/dt=\mu (t,Y)-\frac{1}{2}\mathrm{\nabla }\log p_t(Y)$ started from $p_0$. At any time $t>0$, the marginal distributions of $X_t$ and $Y_t$ both exactly equal $p_t(dx)$.

The diffusion-ODE trick: application to DDPM

Consider a DDPM with forward dynamics given by an Ornstein-Uhlenbeck (OU) process

$$d\overrightarrow{X}=-\frac{1}{2}\overrightarrow{X}dt+dB$$

and initial condition $X_0\sim p_0\equiv {\pi }_{\mathrm{data}}$. As explained in these notes, it is relatively straightforward to estimate the score function

$$\mathcal{S}(t,x_t)=\mathrm{\nabla }\log p_t(x_t)$$

from data. This means that the forward OU process can be replaced by the forward ODE

$$\frac{d}{dt}\overrightarrow{Y_t}=-\frac{1}{2}\overrightarrow{Y_t}-\frac{1}{2}\mathcal{S}(t,\overrightarrow{Y_t})=F(t,Y_t)$$

with $F(t,x)=-\frac{1}{2}x-\frac{1}{2}\mathcal{S}(t,x)$. Similarly, the reverse diffusion (i.e. the “denoising” diffusion) defined as ${\overleftarrow{X}}_s=X_{T-s}$ follows the dynamics

$$d{\overleftarrow{X}}_s=\frac{1}{2}{\overleftarrow{X}}_{s}ds+\mathcal{S}(T-s,{\overleftarrow{X}}_s)+dW.$$

As described for the first time in the beautiful article (Song et al. 2020), the diffusion-ODE trick now shows that the denoising diffusion can be replaced by a denoising ODE with dynamics

$$\begin{array}{rl}\frac{d}{ds}\overleftarrow{Y_s}& =\frac{1}{2}\overleftarrow{Y_s}ds+\mathcal{S}(T-s,\overleftarrow{Y_s})-\frac{1}{2}\mathcal{S}(T-s,{\overleftarrow{X}}_s)\\ & =\frac{1}{2}\overleftarrow{Y_s}ds+\frac{1}{2}\mathcal{S}(T-s,\overleftarrow{Y_s})=-F(T-s,\overleftarrow{Y_s}).\end{array}$$

Interestingly [and I do not know whether there was an obvious way of seeing this from the start], this shows that the forward and backward ODE are actually the same but run forward and backward in time. They corresponds to the ODE described by the vector field

$$\begin{array}{}\text{(2)}& F(t,x)=-\frac{1}{2}x-\frac{1}{2}\mathcal{S}(t,x).\end{array}$$

The animation belows display the denoising ODE and the associated vector field Equation 2.

Likelihood computation

With the diffusion-ODE trick, we have just seen that it is possible to build a vector fields $F[0,T]\times {\mathbb{R}}^d\to {\mathbb{R}}^d$ such that the forward ODE

$$\begin{array}{}\text{(3)}& \frac{d}{dt}\overrightarrow{Y_t}=F(t,\overrightarrow{Y_t})\text{initialized at}{\overrightarrow{Y}}_0\sim {\pi }_{\mathrm{data}}\end{array}$$

and the backward ODE defined as

$$\frac{d}{ds}\overleftarrow{Y_s}=-F(T-s,\overleftarrow{Y_s})\text{initialized at}{\overleftarrow{Y}}_0\sim {\pi }_{\mathrm{ref}}$$

are such that ${\overrightarrow{Y}}_T\approx {\pi }_{\mathrm{ref}}$ and ${\overleftarrow{Y}}_T\approx {\pi }_{\mathrm{data}}$.

In general, consider a vector field $F(t,x)$ and a bunch of particles distributed according to a distribution $p_t$ at time $t$. If each particle follows the vector field for an amount of time $\delta \ll 1$, the particles that were in the vicinity of some $x\in {\mathbb{R}}^d$ at time $t$ end up in the vicinity of $x+F(x,t)\delta$ at time $t+\delta$. At the same time, a volume element $dx$ around $x\in {\mathbb{R}}^d$ gets stretch by a factor $1+\delta \mathrm{Tr}[\mathrm{Jac}F(x,t)]=1+\delta \mathrm{div}F(x,t)$ while following the vector field $F$, which means that the density of particles at time $t+\delta$ and around $x+F(x,t)\delta$ equals $p_t(x)/[1+\delta \mathrm{div}F(x,t)]$. In other words $\log p_{t+\delta }(x+F(x,t)\delta )\approx \log p_t(x)-\delta \mathrm{div}F(x,t)$. This means that if we follows a trajectory of $\frac{d}{dt}{X}_t=F(t,X_t)$ one gets

$$\log p_T(X_T)=\log p_0(X_0)-{\int }_{t=0}^T\mathrm{div}F(X_t,t)dt.$$

That is the Lagrangian description of the density $p_t$ of particles. Indeed, one could directly get this identity by differentiating $p_t(X_t)$ with respect to time while using the continuity Equation 1. When applied to the DDPM, this gives a way to assign likelihood the data samples, namely

$$\log {\pi }_{\mathrm{data}}(x)=\log {\pi }_{\mathrm{ref}}(\overrightarrow{Y_T})+{\int }_{t=0}^T\mathrm{div}F(t,\overrightarrow{Y_t})dt$$

where $\overrightarrow{Y_t}$ is trajectory of the forward ODE Equation 3 initialized as $\overrightarrow{Y_0}=x$. Note that in high-dimensional setting, it may be computationally expensive to compute the divergence term $\mathrm{div}F(t,\overrightarrow{Y_t})$ since it typically is $d$ times slower that a gradient computation; for this reason, it is often advocated to use the Hutchinson trace estimator to get an unbiased estimate of it at a much lower computational cost.

References

Song, Yang, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2020. “Score-Based Generative Modeling Through Stochastic Differential Equations.” ICLR 2021. https://arxiv.org/abs/2011.13456.