Denoising Diffusion Probabilistic Models (DDPM)

This is a PyTorch implementation/tutorial of the paper Denoising Diffusion Probabilistic Models.

In simple terms, we get an image from data and add noise step by step. Then We train a model to predict that noise at each step and use the model to generate images.

The following definitions and derivations show how this works. For details please refer to the paper.

Forward Process

The forward process adds noise to the data $x_0\sim q(x_0)$, for $T$ timesteps.

where ${\beta }_1,\dots ,{\beta }_T$ is the variance schedule.

We can sample $x_t$ at any timestep $t$ with,

where ${\alpha }_t=1-{\beta }_t$ and $\overline{{\alpha }_t}={\prod }_{s=1}^t{\alpha }_s$

Reverse Process

The reverse process removes noise starting at $p(x_T)=\mathcal{N}(x_T;0,\mathbf{I})$ for $T$ time steps.

$\theta$ are the parameters we train.

Loss

We optimize the ELBO (from Jenson's inequality) on the negative log likelihood.

The loss can be rewritten as follows.

$D_{KL}(q(x_T|x_0)\Vert p(x_T))$ is constant since we keep ${\beta }_1,\dots ,{\beta }_T$ constant.

Computing $L_{t-1}=D_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))$

The forward process posterior conditioned by $x_0$ is,

The paper sets ${\Sigma }_{\theta }(x_t,t)={\sigma }_t^2\mathbf{I}$ where ${\sigma }_t^2$ is set to constants ${\beta }_t$ or $\widetilde{{\beta }_t}$.

Then, $$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t^2\mathbf{I})$$

For given noise $ϵ\sim \mathcal{N}(0,\mathbf{I})$ using $q(x_t|x_0)$

This gives,

Re-parameterizing with a model to predict noise

where ${ϵ}_{\theta }$ is a learned function that predicts $ϵ$ given $(x_t,t)$.

This gives,

That is, we are training to predict the noise.

Simplified loss

$$L_{\text{simple}}(\theta )={\mathbb{E}}_{t,x_0,ϵ}[\Vert ϵ-{ϵ}_{\theta }(\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}ϵ,t){\Vert }^2]$$

This minimizes $-\log p_{\theta }(x_0|x_1)$ when $t=1$ and $L_{t-1}$ for $t>1$ discarding the weighting in $L_{t-1}$. Discarding the weights $\frac{{{\beta }_t}^2}{2{\sigma }_t^2{\alpha }_t(1-\overline{{\alpha }_t})}$ increase the weight given to higher $t$ (which have higher noise levels), therefore increasing the sample quality.

This file implements the loss calculation and a basic sampling method that we use to generate images during training.

Here is the UNet model that gives ${ϵ}_{\theta }(x_t,t)$ and training code. This file can generate samples and interpolations from a trained model.

Denoise Diffusion

• eps_model is ${ϵ}_{\theta }(x_t,t)$ model
• n_steps is $t$
• device is the device to place constants on

Get $q(x_t|x_0)$ distribution

Sample from $q(x_t|x_0)$

Sample from $p_{\theta }(x_{t-1}|x_t)$

Simplified Loss

$$L_{\text{simple}}(\theta )={\mathbb{E}}_{t,x_0,ϵ}[\Vert ϵ-{ϵ}_{\theta }(\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}ϵ,t){\Vert }^2]$$