Project 5
Diffusion Models: Sampling, Inpainting, and Optical Illusions
Diffusion Models: Sampling, Inpainting, and Optical Illusions
This project explores diffusion models using the DeepFloyd IF model. We implement diffusion sampling loops, experiment with text-to-image generation, and extend the model to solve various tasks including inpainting, image editing, visual anagrams, and hybrid images.
Random Seed: 42 (used consistently throughout all experiments)
Part A: Diffusion Models with Pre-trained Models
This part explores diffusion models using the DeepFloyd IF model. We implement diffusion sampling loops, experiment with text-to-image generation, and extend the model to solve various tasks including inpainting, image editing, visual anagrams, and hybrid images.
Text-to-Image Generation
We experimented with different text prompts and generated images using various num_inference_steps values. The model successfully captures the essence of text prompts, with quality improving significantly as we increase the number of inference steps.
20 Steps:
"a man wearing a hat"
"a lithograph of a skull"
"a high quality photo"
50 Steps:
"a man wearing a hat"
"a lithograph of a skull"
"a high quality photo"
Forward Process and Denoising
1.1 Implementing the Forward Process
A key part of diffusion is the forward process, which takes a clean image and adds noise to it. The forward process is defined by:
\[q(\mathbf{x}_t | \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t) \mathbf{I})\]which is equivalent to computing:
\[\mathbf{x}_t = \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}, \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\]That is, given a clean image $ \mathbf{x}_0 $, we get a noisy image $ \mathbf{x}_t $ at timestep $ t $ by sampling from a Gaussian with mean $ \sqrt{\bar{\alpha}_t} \mathbf{x}_0 $ and variance $ (1 - \bar{\alpha}_t) \mathbf{I} $. Note that the forward process is not just adding noise—we also scale the image. Here $ \bar{\alpha}_t $ is the cumulative product of noise coefficients, where $ t=0 $ corresponds to a clean image, and larger $ t $ corresponds to more noise. Thus, $ \bar{\alpha}_t $ is close to 1 for small $ t $, and close to 0 for large $ t $.
We implemented the forward(im, t) function and applied it to the Berkeley Campanile test image (resized to 64×64) at different noise levels:
As expected, the images become progressively noisier with increasing timestep values.
1.2 Classical Denoising
Let’s try to denoise these images using classical methods. We take noisy images for timesteps [250, 500, 750] and use Gaussian blur filtering to try to remove the noise. Getting good results should be quite difficult, if not impossible.
Gaussian blur removes some noise but also blurs important image details, making it ineffective for heavy noise levels. This demonstrates why we need learned denoising models.
1.3 One-Step Denoising
Now, we’ll use a pretrained diffusion model to denoise. The actual denoiser can be found at stage_1.unet. This is a UNet that has already been trained on a very large dataset of $ (\mathbf{x}_t, t) $ pairs of images. We can use it to recover Gaussian noise from the image. Then, we can remove this noise to recover (something close to) the original image.
For the 3 noisy images from 1.2 (t = [250, 500, 750]):
- Use the forward function to add noise to the Campanile
- Estimate the noise in the noisy image by passing it through
stage_1.unet - Remove the noise from the noisy image to obtain an estimate of the original image
The denoising process works as follows: given a noisy image $x_t$ at timestep $t$, the UNet predicts the noise $\epsilon_\theta(x_t, t)$. From this noise estimate, we can recover an estimate of the clean image:
\[\hat{\mathbf{x}}_0 = \frac{\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}_\theta}{\sqrt{\bar{\alpha}_t}}\] 
The diffusion model performs much better than Gaussian blur, successfully projecting the image onto the natural image manifold. However, quality degrades with more noise, which makes sense as the problem is much harder with more noise!
1.4 Iterative Denoising
Diffusion models are designed to denoise iteratively. To speed up the process, we use strided timesteps starting at 990 with step size 30, reaching 0. On each step, we transition from $ x_t $ to $ x_{t-1} $ using:
\[\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \right) + \sigma_t \mathbf{z}\]where $ \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) $ is the noise predicted by the UNet, and $ \mathbf{z} $ is predicted noise added via add_variance with variance $ \sigma_t $.
We add noise to the test image at strided_timesteps[10] and run iterative_denoise with i_start = 10. Displaying every 5th step:
Iterative Denoising Process:
Comparison with Other Methods:
Iterative denoising produces significantly better results than one-step denoising, gradually removing noise while preserving image structure.
1.5 Diffusion Model Sampling
In part 1.4, we used the diffusion model to denoise an image. Another thing we can do with the iterative_denoise function is to generate images from scratch. We can do this by setting i_start = 0 and passing im_noisy as random noise. This effectively denoises pure noise.
We generate images from scratch by denoising pure noise (starting at i_start = 0). Here are 5 samples using the prompt “a high quality photo”:
The quality of the images will not be spectacular, but should be reasonable images. We will fix this in the next section with CFG.
1.6 Classifier-Free Guidance (CFG)
CFG dramatically improves image quality by combining conditional and unconditional noise estimates. We used a CFG scale of 7 for all subsequent experiments.
1.7 Image-to-Image Translation (SDEdit)
SDEdit projects images onto the natural image manifold by adding noise and then denoising. The amount of noise controls the strength of the edit.
Web Images
Web Image:
Campanile SDEdit:
Hand-Drawn Images
SDEdit works particularly well for non-realistic inputs, transforming sketches and paintings into photorealistic results.
Hand-Drawn Image 1:
Hand-Drawn Image 2:
1.8 Inpainting
We implemented inpainting using the RePaint algorithm, which preserves original image pixels outside the mask while denoising the masked region.
Inpainting Example 1:
Inpainting Example 2:
Inpainting Example 3 (Campanile):
1.9 Text-Conditional Image Editing
We extended SDEdit to use text prompts for guided image editing. The prompt guides the transformation while preserving the original structure.
Text-Conditional Edit 1:
Prompt: “a rocket ship”
Text-Conditional Edit 2:
Prompt: “a photo of a dog”
Text-Conditional Edit 3:
Prompt: “a photo of a dog”
1.10 Visual Anagrams
Visual anagrams create images that reveal different content when flipped upside down. We achieve this by averaging noise estimates from the original and flipped images with different prompts.
Visual Anagram 1:
Prompt: "a photo of a dog"
Prompt: "a photo of a man"
Visual Anagram 2:
Prompt: "a photo of a dog"
Prompt: "a rocket ship"
Visual Anagram 3:
Prompt: "an oil painting of an old man"
Prompt: "an oil painting of people around a campfire"
1.11 Hybrid Images
We implemented hybrid images using Factorized Diffusion, combining low frequencies from one prompt with high frequencies from another using Gaussian blur (kernel size 33, sigma 2).
Prompts: "a photo of a dog" + "a photo of a man"
Prompts: "a rocket ship" + "a photo of a dog"
Prompts: "a lithograph of a skull" + "a lithograph of waterfalls"
Part B: Training a Flow Matching Model on MNIST
In this part, we train our own flow matching model on the MNIST dataset. We start with a simple one-step denoiser and progressively build up to a full flow matching model with time and class conditioning.
B.1: Training a Single-Step Denoising UNet
We begin by building a simple one-step denoiser. Given a noisy image $x_{\text{noisy}}$, we aim to train a denoiser $f_\theta$ such that it maps $x_{\text{noisy}}$ to a clean image $x_0$. We optimize over an L2 loss:
\[\mathcal{L} = \|f_\theta(x_{\text{noisy}}) - x_0\|^2\]1.1 Implementing the UNet
The denoiser is implemented as a UNet, consisting of downsampling and upsampling blocks with skip connections. The architecture uses standard operations including Conv2d, ConvTranspose2d, BatchNorm2d, GELU activation, and AvgPool2d.
Key Components:
- Conv: Convolutional layer that doesn’t change image resolution, only channel dimension
- DownConv: Convolutional layer that downsamples by 2
- UpConv: Convolutional layer that upsamples by 2
- Flatten: Average pooling that flattens a 7×7 tensor into 1×1
- Unflatten: Convolutional layer that unflattens/upsamples 1×1 tensor into 7×7
- Concat: Channel-wise concatenation using
torch.cat()
The UNet uses hidden dimension $D = 128$ as a hyperparameter.
1.2 Using the UNet to Train a Denoiser
To train our denoiser, we generate training data pairs $(x_{\text{noisy}}, x_0)$ where each $x_0$ is a clean MNIST digit. For each training batch, we generate $x_{\text{noisy}}$ from $x_0$ using the noising process:
\[x_{\text{noisy}} = x_0 + \sigma \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)\]where $\sigma$ is the noise level.
Visualization of the Noising Process:
We visualize the different noising processes over $\sigma$, assuming normalized $x_0$. As $\sigma$ increases, the images become progressively noisier.
Clean MNIST digits with increasing noise levels $\sigma$. Images become progressively noisier as $\sigma$ increases.
1.2.1 Training
We train the model with the following setup:
- Dataset: MNIST training set (shuffled)
- Batch size: 256
- Epochs: 5
- Model: UNet with $D = 128$
- Optimizer: Adam with learning rate $10^{-4}$
- Noise level: $\sigma = 0.5$
The model is trained to denoise images with noise level $\sigma = 0.5$.
Loss decreases over 5 epochs of training
1.2.2 Out-of-Distribution Testing
We test the denoiser on different noise levels $\sigma$ that it wasn’t trained on. This evaluates the model’s generalization capability.
Denoising results with varying noise levels $\sigma$ on the same test image
1.2.3 Denoising Pure Noise
To explore generative capabilities, we train the model to denoise pure random Gaussian noise. This is equivalent to starting with a blank canvas $x_{\text{noisy}} = \epsilon$ where $\epsilon \sim \mathcal{N}(0, I)$ and denoising it to get a clean image $x_0$.
The model learns to predict the centroid of the training distribution when denoising pure noise. With MSE loss, the model minimizes the sum of squared distances to all training examples, resulting in blurry average-like outputs that don't resemble specific digits.
Observations:
- The generated outputs appear as blurry, averaged versions of digits
- They don’t resemble specific digits (0-9) but rather a centroid of the training distribution
- This happens because with MSE loss, the model learns to predict the point that minimizes the sum of squared distances to all training examples, which is the centroid in the context of clustering
B.2: Training a Flow Matching Model
One-step denoising doesn’t work well for generative tasks. Instead, we need iterative denoising using flow matching. We train a UNet model to predict the “flow” from noisy data to clean data. In our flow matching setup, we sample a pure noise image $x_1$ and generate a realistic image $x_0$.
For iterative denoising, we define intermediate noisy samples using linear interpolation between noisy $x_1$ and clean $x_0$:
\[x_t = (1-t) \cdot x_1 + t \cdot x_0, \quad t \in [0, 1]\]This is a vector field describing the position of a point $x_t$ at time $t$ relative to the clean data distribution $x_0$ and the noisy data distribution $x_1$. The flow (velocity) is:
\[v_t = \frac{dx_t}{dt} = x_0 - x_1\]Our aim is to learn a UNet $f_\theta(x_t, t)$ which approximates this flow $v_t$, giving us the learning objective:
\[\mathcal{L} = \|f_\theta(x_t, t) - (x_0 - x_1)\|^2\]2.1 Adding Time Conditioning to UNet
We inject the scalar time $t$ into our UNet model using FCBlock (fully-connected block) operations. The conditioning signal $t$ is normalized to be in the range $[0, 1]$ and embedded into the network at two locations:
- Modulated Unflatten: The unflatten operation is multiplied by the first time embedding $t_1$
- Modulated Up1: The first upsampling block is multiplied by the second time embedding $t_2$
This allows the model to condition its predictions on the current timestep $t$.
2.2 Training the UNet
We train the time-conditioned UNet with the following setup:
- Dataset: MNIST training set (shuffled)
- Batch size: 64
- Model: Time-conditioned UNet with $D = 64$
- Optimizer: Adam with initial learning rate $10^{-2}$
- Scheduler: Exponential learning rate decay with $\gamma = 0.95$ (applied after each epoch)
During training, we pick a random image $x_0$ from the training set, a random timestep $t \sim \mathcal{U}(0, 1)$, add noise to get $x_t = (1-t) \cdot x_1 + t \cdot x_0$ where $x_1 \sim \mathcal{N}(0, I)$, and train the model to predict the flow at $t$.
Loss decreases over the training process with exponential learning rate decay
2.3 Sampling from the UNet
We use iterative denoising to sample from the trained model. Starting from pure noise $x_1 \sim \mathcal{N}(0, I)$, we iteratively apply the predicted flow to generate clean images.
Sampling Algorithm:
- Initialize $x_1 \sim \mathcal{N}(0, I)$
- For $t$ from 1 to 0 (in steps of $\Delta t$):
- Predict flow: $v_t = f_\theta(x_t, t)$
- Update: $x_{t-\Delta t} = x_t - \Delta t \cdot v_t$
- Return $x_0$
Legible digits emerge through iterative denoising (epochs 1, 5, 10)
2.4 Adding Class-Conditioning to UNet
To improve results and provide control over image generation, we add class conditioning to the UNet. The class conditioning vector $c$ is a one-hot vector representing digits 0-9. We implement classifier-free guidance by dropping the class conditioning 10% of the time ($p = 0.1$) during training.
The conditioning is applied at two locations:
- Modulated Unflatten: $c_1 \cdot \text{unflatten} + t_1$
- Modulated Up1: $c_2 \cdot \text{up1} + t_2$
This allows the model to condition on both time $t$ and class $c$.
2.5 Training the Class-Conditioned UNet
Training follows the same procedure as time-only conditioning, with the addition of class conditioning and periodic unconditional generation (10% dropout).
Class conditioning enables faster convergence, requiring only 10 epochs
2.6 Sampling from the Class-Conditioned UNet
We sample with class-conditioning and use classifier-free guidance with guidance scale $w_{\text{cfg}} = 5.0$. The sampling algorithm combines conditional and unconditional predictions:
\[\hat{v}_t = v_{\text{uncond}} + w_{\text{cfg}} \cdot (v_{\text{cond}} - v_{\text{uncond}})\]where $v_{\text{cond}} = f_\theta(x_t, t, c)$ and $v_{\text{uncond}} = f_\theta(x_t, t, 0)$.
Generating 4 instances of each digit (0-9) with class control (epochs 1, 5, 10)
Summary
This project explored diffusion models from two perspectives: using pre-trained models for various tasks and training our own flow matching model on MNIST.
Part A (Pre-trained Models): We worked with the DeepFloyd IF diffusion model to implement sampling loops, text-to-image generation, and extended applications including inpainting, image editing, visual anagrams, and hybrid images. Key techniques included iterative denoising, classifier-free guidance (CFG), SDEdit for image-to-image translation, and creative applications like visual anagrams that reveal different content when flipped.
Part B (Training from Scratch): We built and trained UNet-based models on MNIST, starting with simple one-step denoising and progressing to full flow matching models. The project demonstrated that one-step denoising fails for generative tasks (producing blurry centroids), while iterative flow matching successfully generates realistic digits. Adding time and class conditioning improved generation quality and provided control over the output. Classifier-free guidance further enhanced the results by combining conditional and unconditional predictions.
The project successfully demonstrated both the practical applications of diffusion models and the fundamental principles behind training them, showing how iterative denoising with proper conditioning can generate high-quality images from noise.