The Annotated Diffusion Model

Back to Articles The Annotated Diffusion Model Published June 7, 2022 Update on GitHub Upvote 365 +359 Niels Rogge nielsr Follow Kashif Rasul kashif Follow What is a diffusion model? In more mathematical form Defining an objective function (by reparametrizing the mean) The neural network Network...

Back to Articles The Annotated Diffusion Model Published June 7, 2022 Update on GitHub Upvote 365 +359 Niels Rogge nielsr Follow Kashif Rasul kashif Follow What is a diffusion model? In more mathematical form Defining an objective function (by reparametrizing the mean) The neural network Network helpers Position embeddings ResNet block Attention module Group normalization Conditional U-Net Defining the forward diffusion process Define a PyTorch Dataset + DataLoader Sampling Train the model Sampling (inference) Follow-up reads In this blog post, we'll take a deeper look into Denoising Diffusion Probabilistic Models (also known as DDPMs, diffusion models, score-based generative models or simply autoencoders) as researchers have been able to achieve remarkable results with them for (un)conditional image/audio/video generation. Popular examples (at the time of writing) include GLIDE and DALL-E 2 by OpenAI, Latent Diffusion by the University of Heidelberg and ImageGen by Google Brain. We'll go over the original DDPM paper by (Ho et al., 2020), implementing it step-by-step in PyTorch, based on Phil Wang's implementation - which itself is based on the original TensorFlow implementation. Note that the idea of diffusion for generative modeling was actually already introduced in (Sohl-Dickstein et al., 2015). However, it took until (Song et al., 2019) (at Stanford University), and then (Ho et al., 2020) (at Google Brain) who independently improved the approach. Note that there are several perspectives on diffusion models. Here, we employ the discrete-time (latent variable model) perspective, but be sure to check out the other perspectives as well. Alright, let's dive in! from IPython.display import Image Image(filename='assets/78_annotated-diffusion/ddpm_paper.png') We'll install and import the required libraries first (assuming you have PyTorch installed). !pip install -q -U einops datasets matplotlib tqdm import math from inspect import isfunction from functools import partial %matplotlib inline import matplotlib.pyplot as plt from tqdm.auto import tqdm from einops import rearrange, reduce from einops.layers.torch import Rearrange import torch from torch import nn, einsum import torch.nn.functional as F What is a diffusion model? A (denoising) diffusion model isn't that complex if you compare it to other generative models such as Normalizing Flows, GANs or VAEs: they all convert noise from some simple distribution to a data sample. This is also the case here where a neural network learns to gradually denoise data starting from pure noise. In a bit more detail for images, the set-up consists of 2 processes: a fixed (or predefined) forward diffusion process qqq of our choosing, that gradually adds Gaussian noise to an image, until you end up with pure noise a learned reverse denoising diffusion process pθp_\thetapθ​, where a neural network is trained to gradually denoise an image starting from pure noise, until you end up with an actual image. Both the forward and reverse process indexed by ttt happen for some number of finite time steps TTT (the DDPM authors use T=1000T=1000T=1000). You start with t=0t=0t=0 where you sample a real image x0\mathbf{x}_0x0​ from your data distribution (let's say an image of a cat from ImageNet), and the forward process samples some noise from a Gaussian distribution at each time step ttt, which is added to the image of the previous time step. Given a sufficiently large TTT and a well behaved schedule for adding noise at each time step, you end up with what is called an isotropic Gaussian distribution at t=Tt=Tt=T via a gradual process. In more mathematical form Let's write this down more formally, as ultimately we need a tractable loss function which our neural network needs to optimize. Let q(x0)q(\mathbf{x}_0)q(x0​) be the real data distribution, say of "real images". We can sample from this distribution to get an image, x0∼q(x0)\mathbf{x}_0 \sim q(\mathbf{x}_0)x0​∼q(x0​). We define the forward diffusion process q(xt∣xt−1)q(\mathbf{x}_t | \mathbf{x}_{t-1})q(xt​∣xt−1​) which adds Gaussian noise at each time step ttt, according to a known variance schedule 0 0: arr.append(remainder) return arr results_folder = Path("./results") results_folder.mkdir(exist_ok = True) save_and_sample_every = 1000 Below, we define the model, and move it to the GPU. We also define a standard optimizer (Adam). from torch.optim import Adam device = "cuda" if torch.cuda.is_available() else "cpu" model = Unet( dim=image_size, channels=channels, dim_mults=(1, 2, 4,) ) model.to(device) optimizer = Adam(model.parameters(), lr=1e-3) Let's start training! from torchvision.utils import save_image epochs = 6 for epoch in range(epochs): for step, batch in enumerate(dataloader): optimizer.zero_grad() batch_size = batch["pixel_values"].shape[0] batch = batch["pixel_values"].to(device) # Algorithm 1 line 3: sample t uniformally for every example in the batch t = torch.randint(0, timesteps, (batch_size,), device=device).long() loss = p_losses(model, batch, t, loss_type="huber") if step % 100 == 0: print("Loss:", loss.item()) loss.backward() optimizer.step() # save generated images if step != 0 and step % save_and_sample_every == 0: milestone = step // save_and_sample_every batches = num_to_groups(4, batch_size) all_images_list = list(map(lambda n: sample(model, batch_size=n, channels=channels), batches)) all_images = torch.cat(all_images_list, dim=0) all_images = (all_images + 1) * 0.5 save_image(all_images, str(results_folder / f'sample-{milestone}.png'), nrow = 6) Output: ---------------------------------------------------------------------------------------------------- Loss: 0.46477368474006653 Loss: 0.12143351882696152 Loss: 0.08106148988008499 Loss: 0.0801810547709465 Loss: 0.06122320517897606 Loss: 0.06310459971427917 Loss: 0.05681884288787842 Loss: 0.05729678273200989 Loss: 0.05497899278998375 Loss: 0.04439849033951759 Loss: 0.05415581166744232 Loss: 0.06020551547408104 Loss: 0.046830907464027405 Loss: 0.051029372960329056 Loss: 0.0478244312107563 Loss: 0.046767622232437134 Loss: 0.04305662214756012 Loss: 0.05216279625892639 Loss: 0.04748568311333656 Loss: 0.05107741802930832 Loss: 0.04588869959115982 Loss: 0.043014321476221085 Loss: 0.046371955424547195 Loss: 0.04952816292643547 Loss: 0.04472338408231735 Sampling (inference) To sample from the model, we can just use our sample function defined above: # sample 64 images samples = sample(model, image_size=image_size, batch_size=64, channels=channels) # show a random one random_index = 5 plt.imshow(samples[-1][random_index].reshape(image_size, image_size, channels), cmap="gray") Seems like the model is capable of generating a nice T-shirt! Keep in mind that the dataset we trained on is pretty low-resolution (28x28). We can also create a gif of the denoising process: import matplotlib.animation as animation random_index = 53 fig = plt.figure() ims = [] for i in range(timesteps): im = plt.imshow(samples[i][random_index].reshape(image_size, image_size, channels), cmap="gray", animated=True) ims.append([im]) animate = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) animate.save('diffusion.gif') plt.show() Follow-up reads Note that the DDPM paper showed that diffusion models are a promising direction for (un)conditional image generation. This has since then (immensely) been improved, most notably for text-conditional image generation. Below, we list some important (but far from exhaustive) follow-up works: Improved Denoising Diffusion Probabilistic Models (Nichol et al., 2021): finds that learning the variance of the conditional distribution (besides the mean) helps in improving performance Cascaded Diffusion Models for High Fidelity Image Generation (Ho et al., 2021): introduces cascaded diffusion, which comprises a pipeline of multiple diffusion models that generate images of increasing resolution for high-fidelity image synthesis Diffusion Models Beat GANs on Image Synthesis (Dhariwal et al., 2021): show that diffusion models can achieve image sample quality superior to the current state-of-the-art generative models by improving the U-Net architecture, as well as introducing classifier guidance Classifier-Free Diffusion Guidance (Ho et al., 2021): shows that you don't need a classifier for guiding a diffusion model by jointly training a conditional and an unconditional diffusion model with a single neural network Hierarchical Text-Conditional Image Generation with CLIP Latents (DALL-E 2) (Ramesh et al., 2022): uses a prior to turn a text caption into a CLIP image embedding, after which a diffusion model decodes it into an image Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding (ImageGen) (Saharia et al., 2022): shows that combining a large pre-trained language model (e.g. T5) with cascaded diffusion works well for text-to-image synthesis Note that this list only includes important works until the time of writing, which is June 7th, 2022. For now, it seems that the main (perhaps only) disadvantage of diffusion models is that they require multiple forward passes to generate an image (which is not the case for generative models like GANs). However, there's research going on that enables high-fidelity generation in as few as 10 denoising steps. More Articles from our Blog guidecollaborationdiffusers LoRA training scripts of the world, unite! 80 January 2, 2024 guidediffusionstable-diffusion Train your ControlNet with diffusers 37 March 24, 2023 Community yangrk Mar 13, 2025 The generate images is randomly for different categories? Since the dataset contains different categories of images See translation 1 reply · kashif Article author Mar 13, 2025 right because this model is being trained in an unconditional fashion... so the samples will be random across the 10 classes See translation 👍 1 1 + hsz19 Jun 10, 2025 Your article is excellent and has been incredibly helpful to me. I'm wondering why is dim=image_size used when calling Unet? It seems that dim affects the number of output channels in the network, so why is it assigned the value of image size? See translation 👍 2 2 + Reply EditPreview Upload images, audio, and videos by dragging in the text input, pasting, or clicking here. Tap or paste here to upload images Comment · Sign up or log in to comment Upvote 365 +353

Source: Hugging Face — Published June 7, 2022 at 00:00 UTC — Category: Models