Implementation of the Denoising Diffusion Probabilistic Model (DDPM)#

Numerical Implementation of the Forward Process#

To numerically simulate the forward diffusion process, we divide the time range $[0,T]$ into intervals of length ${\beta }_i$, where $i=1,\dots ,n$. We denote the intermediate times as $t_i={\sum }_{j=0}^{i-1}{\beta }_j$.

The vanilla discretization of the $\text{Forward Process}$ is given by:

where we approximate $d{\mathbf{W}}_t$ as $\sqrt{dt}{\mathbf{ϵ}}_i$, ${\mathbf{ϵ}}_i$ is standard Gaussian random variable. (refer to the previous section).

A more subtle but equivalent implementation is the variance-preserving (VP) form ¹:

This formulation ensures that if ${\mathbf{x}}_0$ is initialized with unit variance, then the variance of ${\mathbf{x}}_i$ remains equal to 1. It gradually adds a small amount of Gaussian noise to the image at each time step $i$, gradually contaminating the image until ${\mathbf{x}}_n\sim \mathcal{N}(0,I)$.

WARNING

Note that our interpretation of $\beta$ differs from that in ¹, treating $\beta$ as a varying time-step size to solve the autonomous SDE (1.5 OU process noise) instead of a time-dependent SDE. Our interpretation greatly simplifies future analysis, but it holds only if every ${\beta }_i$ is sufficiently small.

Instead of expressing the iterative relationship between ${\mathbf{x}}_i$ and ${\mathbf{x}}_{i-1}$, we can directly represent the dependency of ${\mathbf{x}}_i$ on ${\mathbf{x}}_0$ using the following forward relation:

where ${\bar{\alpha }}_i={\prod }_{j=0}^{i-1}(1-{\beta }_j)$ denotes the contamination weight, and ${\overline{\mathbf{ϵ}}}_i$ represents standard Gaussian noise.

TIP

An useful property we shall exploit later is that for infinitesimal time steps $\beta$, the contamination weight ${\bar{\alpha }}_i$ is the exponential of the diffusion time $t_i$

$$\underset{{\max }_j{\beta }_j\to 0}{\lim }{\bar{\alpha }}_i\to e^{-t_i}.$$

Numerical Implementation of the Backward Process#

The backward diffusion process is used to sample from the DDPM by removing the noise of an image step by step. It is the time reversed version of the OU process, starting at $x_{0^{\prime }}\sim \mathcal{N}(\mathbf{x}|0,I)$, using the reverse of the OU process (1.5 reverse diffusion process).

The vanilla discretization of the $\text{Backward Process}$ is given by:

where $i^{\prime }=0,\dots ,n$ represents the backward time step, and ${\mathbf{x}}_{i^{\prime }}$ is the image at the $i^{\prime }$th step with time $t_{i^{\prime }}^{\prime }={\sum }_{j=0}^{i^{\prime }-1}{\beta }_{n-1-j}=T-t_{n-i^{\prime }}$.

A more common discretization is:

This formulation is equivalent to the vanilla discretization when ${\beta }_i$ is small. The score function $\mathbf{s}({\mathbf{x}}_{i^{\prime }},T-t_{i^{\prime }}^{\prime })$ is typically modeled by a neural network trained using a denoising objective.

Training the Score Function#

Training the score function requires a training objective. We will show that the score function could be trained with a denoising objective.

DDPM is trained to removes the noise ${\overline{\mathbf{ϵ}}}_i$ from ${\mathbf{x}}_i$ in the forward diffusion process, by training a denoising neural network ${\mathbf{ϵ}}_{\theta }(\mathbf{x},t_i)$ to predict and remove the noise ${\overline{\mathbf{ϵ}}}_i$. This means that DDPM minimizes the denoising objective ²:

where ${\overline{\mathbf{ϵ}}}_i$ is determined from ${\mathbf{x}}_i$ according to the $\text{discrete forward diffusion}$ process.

Now we show that ${\mathbf{ϵ}}_{\theta }$ trained with the above objective is proportional to the score function $\mathbf{s}$. There are two important properties regarding the relationship between the noise ${\overline{\mathbf{ϵ}}}_i$ and the score function $\mathbf{s}$:

1. Gaussian Distribution of ${\mathbf{x}}_i$:
   According to the $\text{discrete forward diffusion}$, the distribution of ${\mathbf{x}}_i$ given ${\mathbf{x}}_0$ is a Gaussian distribution, expressed as:

   $$p({\mathbf{x}}_i|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_i|\sqrt{{\bar{\alpha }}_i}{\mathbf{x}}_0,(1-{\bar{\alpha }}_i)I).$$

2. Proportionality of Noise to Score Function:
   The noise ${\overline{\mathbf{ϵ}}}_i$ is directly proportional to a score function, given by:

   $${\overline{\mathbf{ϵ}}}_i=-\sqrt{1-{\bar{\alpha }}_i}\mathbf{s}({\mathbf{x}}_i|{\mathbf{x}}_0,t_i),$$

   where $\mathbf{s}({\mathbf{x}}_i|{\mathbf{x}}_0,t_i)={\nabla }_{{\mathbf{x}}_i}\log p({\mathbf{x}}_i|{\mathbf{x}}_0)$ represents the score of the conditional probability density $p({\mathbf{x}}_i|{\mathbf{x}}_0)$ at ${\mathbf{x}}_i$.

These properties indicate that the noise ${\overline{\mathbf{ϵ}}}_i$ is directly related to a conditional score function, which connects to the score function $\mathbf{s}(\mathbf{x},t)$ through the above equations.

Now we are very close to our target. The conditional score function $\mathbf{s}({\mathbf{x}}_i|{\mathbf{x}}_0,t_i)$ is connected to the score function $\mathbf{s}(\mathbf{x},t)$ through the following equation:

where $f$ is an arbitrary function and $\mathbf{s}(\mathbf{x},t)={\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})$ is the score function of the probability density of ${\mathbf{x}}_t$.

Substituting the $\text{score-noise relationship}$ into the $\text{denoising objective}$, expanding the squares, and utilizing the above equation, we can derive that the denoising objective is equivalent to a denoising score matching objective:

This objectives says that the denoising neural network ${\mathbf{ϵ}}_{\theta }(\mathbf{x},t_i)$ is trained to approximate a scaled score function $\mathbf{ϵ}(\mathbf{x},t_i)$ ³

Summary:#

We have covered all aspects of the DDPM theory. You can now find a suitable dataset, perform the $\text{discrete forward diffusion}$, train a denoising neural network using the $\text{denoising objective}$, and subsequently generate new samples with the $\text{eps-score relation}$ and the $\text{discrete backward process}$.

What is Next#

In the next section, we will discuss an alternative version of the backward diffusion process: ordinary differential equation (ODE) based backward sampling. This approach serves as the foundation for several modern architectures, such as rectified flow diffusion models.

Stay tuned for the next installment!

Discussion#

If you have questions, suggestions, or ideas to share, please visit the discussion post.