Diffusion Model

The ideas behind diffusion models have been around for many years, but they were first formalized in their modern form in a 2015 paper by Jascha Sohl-Dickstein et al. from Stanford University and UC Berkeley. The authors applied tools from thermodynamics to model a diffusion process, similar to a drop of milk diffusing in a cup of tea. The core idea is to train a model to learn the reverse process: start from the completely mixed state, and gradually “unmix” the milk from the tea. Using this idea, they obtained promising results in image generation, but since GANs produced more convincing images back then, diffusion models did not get as much attention.

Then, in 2020, Jonathan Ho et al., also from UC Berkeley, managed to build a diffusion model capable of generating highly realistic images, which they called a denoising diffusion probabilistic model (DDPM). A few months later, a 2021 paper by OpenAI researchers Alex Nichol and Prafulla Dhariwal analyzed the DDPM architecture and proposed several improvements that allowed DDPMs to finally beat GANs: not only are DDPMs much easier to train than GANs, but the generated images are more diverse and of even higher quality. The main downside of DDPMs, as you will see, is that they take a very long time to generate images, compared to GANs or VAEs.

Diffusion refers to the process of turning a structured signal (an image) into noise step-by-step. By simulating diffusion, we can generate noisy images from our training images, and can train a neural network to try to denoise them. Using the trained network we can simulate the opposite of diffusion, reverse diffusion, which is the process of an image emerging from noise as shown in fig1.

Fig 1. Forward and Reverse(Backward ) Process in Diffusion model

In summary Diffusion models are trained to denoise noisy images, and can generate images by iteratively denoising pure noise.

OK, so let’s start coding! The first thing we need to do is to code the forward process. For this, we will first need to implement the variance schedule

Let’s create a small function to compute α , β , and α̅ , and call it with T = 4,000:

The prepare_batch() function takes a batch of images and adds noise to each of them, using a different random time between 1 and T for each image, and it returns a tuple containing the inputs and the targets:

• The inputs are a dict containing the noisy images and the corresponding times. The function uses equation (4) from the DDPM paper to compute the noisy images in one shot, directly from the original images. It's a shortcut for the forward diffusion process.
• The target is the noise that was used to produce the noisy images.

Now let’s prepare a tf.data.Dataset for training, and one for validation.

As a quick sanity check, let’s take a look at a few training samples, along with the corresponding noise to predict, and the original images (which we get by subtracting the appropriately scaled noise from the appropriately scaled noisy image) in next figures :

Get Redamahmoud’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

Original Images

Fig 2 . Original images of Fashion Minst

Noisy Images

Fig 3. Noisy Images of Fashion Minst

Noisy Images ready for prediction

Fig 3. Predictable Noisy Images of Fashion Minst

Now we’re ready to build the diffusion model itself. It will need to process both images and times. We will encode the times using a sinusoidal encoding, as suggested in the DDPM paper, Given a vector of m integers representing time indices (integers), the layer returns an m × d matrix, where d is the chosen embedding size.

Now let’s build the model. In the Improved DDPM paper, they use a UNet model. We’ll create a UNet-like model, that processes the image through Conv2D + BatchNormalization layers and skip connections, gradually downsampling the image (using MaxPooling layers with strides=2), then growing it back again (using Upsampling2D layers). Skip connections are also added across the downsampling part and the upsampling part. We also add the time encodings to the output of each block, after passing them through a Dense layer to resize them to the right dimension.

Now we are ready to train the model .

After the model trained and got 0.073 val loss , we can use it to generate new images. For this, we just generate Gaussian noise, and pretend this is the result of the diffusion process, and we’re at time T. Then we use the model to predict the image at time T−1, then we call it again to get T−2, and so on, removing a bit of noise at each step. At the end, we get an image that looks like it’s from the Fashion MNIST dataset. The equation for this reverse process is at the top of page 4 in the DDPM paper (step 4 in algorithm 2).

Here’s some of the generation picture result .

Fig 5. Generating images from diffusion model

Thanks for reading

Refences

Hands-On Machine Learning with Scikit-Learn and TensorFlow

Denoising Diffusion Implicit Models on keras Documentaion

Denoising Diffusion Probabilistic Models paper