Fun With Diffusion Models!

I also repeated the same procedure on images from the web and hand-drawn images.

1.7.2. Inpainting

By using a mask to specify regions to edit, we can inpaint images. Here are some examples:

1.7.3. Text-Conditional Image-to-Image Translation

By guiding the denoising process with text prompts, we can create edits that align with the desired description. In this part, I used the text embedding "a childish drawing" to guide the edits. Here are some examples:

1.8. Visual Anagrams

In this part, I created visual anagrams where the image appears as one thing when viewed normally, and another when flipped upside down. Here are some examples:

1.9. Hybrid Images

In this part, I created hybrid images that change appearance based on viewing distance. This uses a similar technique to visual anagrams, but blends high and low frequency components. Here are some examples:

Part B1. Training a Single-Step Denoising UNet

Implementing the UNet

I implemented the UNet architecture as described in the project instructions using PyTorch. The UNet consists of downsampling and upsampling blocks with skip connections. I defined the necessary operations such as Conv, DownConv, UpConv, Flatten, Unflatten, and Concat, and composed them to create a deeper network.

1.2. Using the UNet to Train a Denoiser

I trained the UNet to denoise images from the MNIST dataset. For each training batch, I generated noisy images by adding Gaussian noise with varying levels of noise. The model was optimized using an L2 loss function to minimize the difference between the denoised output and the clean images. I first visualized the noising process on sample MNIST digits:

Then, I trained the UNet for 5 epochs and visualized the denoised results on the test set after the 1st and 5th epochs. I used the model as specified in the project instructions with a hidden dimension of 128 and Adam optimizer with a learning rate of 1e-4. Here are the training loss curve and sample results:

Here are the full denoising results for each digits:

1.2.3. Denoising Pure Noise

I also tested the denoiser on out-of-distribution noise levels and observed how the model performed on different noise levels. Additionally, I trained the model to denoise pure Gaussian noise and visualized the generated outputs after 1 and 5 epochs. Here are the results:

The generated outputs from pure noise exhibited patterns resembling the average of the training images, which is expected due to the nature of the MSE loss function. The model learns to predict the point that minimizes the sum of squared distances to all training examples, leading to outputs that look like the average of the training data distribution. For the higher epochs, the outputs became more refined with smoother edges.

Part B2. Training a Flow Matching Model

2.1. Adding Time Conditioning to UNet

For this part, I modified the UNet architecture to include time conditioning using FCBlocks. The scalar time variable \(t\) was normalized and embedded using two FCBlocks, which were integrated into the Unflatten and UpConv layers of the UNet as described in the project instructions.

2.2. Training the UNet

I trained the time-conditioned UNet on the MNIST dataset to predict the flow from noisy images to clean images at various timesteps. The model was optimized using the Adam optimizer with an initial learning rate of 1e-2 and an exponential learning rate decay scheduler. Here is the training loss curve:

2.3. Sampling from the Time-conditioned UNet

After training, I sampled from the time-conditioned UNet for 1, 5, and 10 epochs. The results showed that legible digits emerged as the number of epochs increased, demonstrating the effectiveness of the flow matching approach. Here are the sampling results:

2.4. Adding Class-Conditioning to UNet

I further modified the UNet to include class conditioning using one-hot vectors for the digit classes (0-9). I implemented dropout to randomly drop the class conditioning vector during training, allowing the model to learn both conditional and unconditional generation.

2.5. Training the Class-Conditioned UNet

I trained the class-conditioned UNet using the same procedure as the time-conditioned UNet, with the addition of class conditioning. The training loss curve is shown below:

2.6. Sampling from the Class-Conditioned UNet

Finally, I sampled from the class-conditioned UNet for 1, 5, and 10 epochs using classifier-free guidance. The results showed that class-conditioning allowed for faster convergence and more accurate digit generation. Here are the sampling results:

Finally, I experimented with training the class-conditioned UNet without the learning rate scheduler. I removed the exponential learning rate scheduler and trained the model using a constant learning rate throughout all epochs. To keep training stable without decay, I lowered the learning rate from \(1\times 10^{-2}\) to \(1\times 10^{-3}\), since a larger fixed step size tended to make the optimization noisy and prevented the loss from settling. With the constant learning rate, the model still converged, but the improvements were more gradual compared to the scheduled run. Below are the sampling results after 10 epochs without the scheduler: