• Conv2d(kernel_size, stride, padding) is nn.Conv2d()
• BN is nn.BatchNorm2d()
• GELU is nn.GELU()
• ConvTranspose2d(kernel_size, stride, padding) is nn.ConvTranspose2d()
• AvgPool(kernel_size) is nn.AvgPool2d()
• D is the number of hidden channels and is a hyperparameter that we will set ourselves.

• (1) Conv is a convolutional layer that doesn't change the image resolution, only the channel dimension.
• (2) DownConv is a convolutional layer that downsamples the tensor by 2.
• (3) UpConv is a convolutional layer that upsamples the tensor by 2.
• (4) Flatten is an average pooling layer that flattens a 7x7 tensor into a 1x1 tensor. 7 is the resulting height and width after the downsampling operations.
• (5) Unflatten is a convolutional layer that unflattens/upsamples a 1x1 tensor into a 7x7 tensor.
• (6) Concat is a channel-wise concatenation between tensors with the same 2D shape. This is simply torch.cat().

We define composed operations using our simple operations in order to make our network deeper. This doesn't change the tensor's height, width, or number of channels, but simply adds more learnable parameters.

• (7) ConvBlock, is similar to Conv but includes an additional Conv. Note that it has the same input and output shape as (1) Conv.
• (8) DownBlock, is similar to DownConv but includes an additional ConvBlock. Note that it has the same input and output shape as (2) DownConv.
• (9) UpBlock, is similar to UpConv but includes an additional ConvBlock. Note that it has the same input and output shape as (3) UpConv.

1.2 Using the UNet to Train a Denoiser

Figure 3: Varying levels of noise on MNIST digits

1.2.1 Training

Now, we will train the model to perform denoising.

• Objective: Train a denoiser to denoise noisy image $z$ with $\sigma =0.5$ applied to a clean image $x$.
• Dataset and dataloader: Use the MNIST dataset via torchvision.datasets.MNIST with flags to access training and test sets. Train only on the training set. Shuffle the dataset before creating the dataloader. Recommended batch size: 256. We'll train over our dataset for 5 epochs.
  • You should only noise the image batches when fetched from the dataloader so that in every epoch the network will see new noised images, improving generalization.
• Model: Use the UNet architecture defined in section 1.1 with recommended hidden dimension D = 128.
• Optimizer: Use Adam optimizer with learning rate of 1e-4.

Figure 4: Training Loss Curve

They should look something like these:

Figure 5: Results on digits from the test set after 1 epoch of training

Figure 6: Results on digits from the test set after 5 epochs of training

1.2.2 Out-of-Distribution Testing

Our denoiser was trained on MNIST digits noised with $\sigma =0.5$. Let's see how the denoiser performs on different $\sigma$'s that it wasn't trained for.

Visualize the denoiser results on test set digits with varying levels of noise $\sigma =[0.0,0.2,0.4,0.5,0.6,0.8,1.0]$.

Figure 7: Results on digits from the test set with varying noise levels.

Deliverables

• A visualization of the noising process using $\sigma =[0.0,0.2,0.4,0.5,0.6,0.8,1.0]$. (figure 3)
• A training loss curve plot every few iterations during the whole training process (figure 4).
• Sample results on the test set after the first and the 5-th epoch (staff solution takes ~7 minutes for 5 epochs on a Colab T4 GPU). (figure 5, 6)
• Sample results on the test set with out-of-distribution noise levels after the model is trained. Keep the same image and vary $\sigma =[0.0,0.2,0.4,0.5,0.6,0.8,1.0]$. (figure 7)

• Since training can take a while, we strongly recommend that you checkpoint your model every epoch onto your personal Google Drive. This is because Colab notebooks aren't persistent such that if you are idle for a while, you will lose connection and your training progress. This consists of:
  • Google Drive mounting.
  • Epoch-wise model & optimizer checkpointing.
  • Model & optimizer resuming from checkpoints.

Part 2: Training a Diffusion Model

Let's revisit the problem we solved in equation B.1:

$$L={\mathbb{E}}_{z,x}\Vert D_{\theta }(z)-x{\Vert }^2.$$

$$\begin{array}{}\text{(B.3)}& L={\mathbb{E}}_{ϵ,z}\Vert {ϵ}_{\theta }(z)-ϵ{\Vert }^2\end{array}$$

where ${ϵ}_{\theta }$ is a UNet trained to predict noise.

For diffusion, we eventually want to sample a pure noise image $ϵ\sim N(0,I)$ and generate a realistic image $x$ from the noise. However, we saw in part A that one-step denoising does not yield good results. Instead, we need to iteratively denoise the image for better results.

Recall in part A that we used equation A.2 to generate noisy images $x_t$ from $x_0$ for some timestep $t$ for $t\in \{0,1,\cdots ,T\}$: $$x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ\text{where}ϵ\sim N(0,1).$$ Intuitively, when $t=0$ we want $x_t$ to be the clean image $x_0$, when $t=T$ we want $x_t$ to be pure noise $ϵ$, and for $t\in \{1,\cdots ,T-1\}$, $x_t$ should be some linear combination of the two. The precise derivation of $\overline{\alpha }$ is beyond the scope of this project (see DDPM paper for more details). Here, we provide you with the DDPM recipe to build a list $\overline{\alpha }$ for $t\in \{0,1,\cdots ,T\}$ utilizing lists $\alpha$ and $\beta$:

• Create a list $\beta$ of length $T$ such that ${\beta }_0=0.0001$ and ${\beta }_T=0.02$ and all other elements ${\beta }_t$ for $t\in \{1,\cdots ,T-1\}$ are evenly spaced between the two.
• ${\alpha }_t=1-{\beta }_t$
• ${\overline{\alpha }}_t=\prod _{s=1}^t{\alpha }_s$ is a cumulative product of ${\alpha }_s$ for $s\in \{1,\cdots ,t\}$.

Now, to denoise image $x_t$, we could simply apply our UNet ${ϵ}_{\theta }$ on $x_t$ and get the noise $ϵ$. However, this won't work very well because the UNet is expecting the noisy image to have a noise variance $\sigma =0.5$ for best results, but the variance of $x_t$ varies with $t$. One could train $T$ separate UNets, but it is much easier to simply condition a single UNet with timestep $t$, giving us our final objective: $$\begin{array}{}\text{(B.4)}& L={\mathbb{E}}_{ϵ,x_0,t}\Vert {ϵ}_{\theta }(x_t,t)-ϵ{\Vert }^2.\end{array}$$

2.1 Adding Time Conditioning to UNet

Figure 8: Conditioned UNet

This uses a new operator called FCBlock (fully-connected block) which we use to inject the conditioning signal into the UNet:

Figure 9: FCBlock for conditioning

Since our conditioning signal $t$ is a scalar, F_in should be of size 1. We also recommend that you normalize $t$ to be in the range [0, 1] before embedding it, i.e. pass in $\frac{t}{T}$.

You can embed $t$ by following this pseudo code:

fc1_t = FCBlock(...)
fc2_t = FCBlock(...)

# the t passed in here should be normalized to be in the range [0, 1]
t1 = fc1_t(t)
t2 = fc2_t(t)

# Follow diagram to get unflatten.
# Replace the original unflatten with modulated unflatten.
unflatten = unflatten + t1
# Follow diagram to get up1.
...
# Replace the original up1 with modulated up1.
up1 = up1 + t2
# Follow diagram to get the output.
...
    

2.2 Training the UNet

Algorithm B.1. Training time-conditioned UNet

• Objective: Train a time-conditioned UNet ${ϵ}_{\theta }(x_t,t)$ to predict the noise in $x_t$ given a noisy image $x_t$ and a timestep $t$.
• Dataset and dataloader: Use the MNIST dataset via torchvision.datasets.MNIST with flags to access training and test sets. Train only on the training set. Shuffle the dataset before creating the dataloader. Recommended batch size: 128. We'll train over our dataset for 20 epochs since this task is more difficult than part A.
  • As shown in algorithm B.1, You should only noise the image batches when fetched from the dataloader.
• Model: Use the time-conditioned UNet architecture defined in section 2.1 with recommended hidden dimension D = 64. Follow the diagram and pseudocode for how to inject the conditioning signal $t$ into the UNet. Remember to normalize $t$ before embedding it.
• Optimizer: Use Adam optimizer with an initial learning rate of 1e-3. We will be using an exponential learning rate decay scheduler with a gamma of ${0.1}^{(1.0/\text{num\_epochs})}$. This can be implemented using scheduler = torch.optim.lr_scheduler.ExponentialLR(...). You should call scheduler.step() after every epoch.

Figure 10: Time-Conditioned UNet training loss curve

2.3 Sampling from the UNet

Algorithm B.2. Sampling from time-conditioned UNet

Deliverables

• A training loss curve plot for the time-conditioned UNet over the whole training process (figure 10).
• Sampling results for the time-conditioned UNet for 5 and 20 epochs.
  • Note: providing a gif is optional.

2.4 Adding Class-Conditioning to UNet

fc1_t = FCBlock(...)
fc1_c = FCBlock(...)
fc2_t = FCBlock(...)
fc2_c = FCBlock(...)

t1 = fc1_t(t)
c1 = fc1_c(c)
t2 = fc2_t(t)
c2 = fc2_c(c)

# Follow diagram to get unflatten.
# Replace the original unflatten with modulated unflatten.
unflatten = c1 * unflatten + t1
# Follow diagram to get up1.
...
# Replace the original up1 with modulated up1.
up1 = c2 * up1 + t1
# Follow diagram to get the output.
...



        

Algorithm B.3. Training class-conditioned UNet

Figure 11: Class-conditioned UNet training loss curve

2.5 Sampling from the Class-Conditioned UNet

Algorithm B.4. Sampling from class-conditioned UNet

Deliverables

• A training loss curve plot for the class-conditioned UNet over the whole training process.
• Sampling results for the class-conditioned UNet for 5 and 20 epochs. Generate 4 instances of each digit as shown above.
  • Note: providing a gif is optional.

Extra Credit

• Improve the UNet Architecture for time-conditional generation
  For ease of explanation and implementation, our UNet architecture above is pretty simple. Modify the UNet (e.g. with skip connections) such that it can fit better during training and sample even better results.
• Implement Rectified Flow
  • Implement rectified flow, which is the state of art diffusion model.
  • You can reference any code on github, but your implementation needs to follow the same code structure as our DDPM implementation.
  • In other words, the code change required should be minimal: only changing the forward and sample functions.

Acknowledgements

This project was a joint effort by Daniel Geng, Ryan Tabrizi, and Hang Gao, advised by Liyue Shen, Andrew Owens, and Alexei Efros.

Last updated 16 April 2025