• 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.

1.2.3 Denoising Pure Noise

To make denoising a generative task, we'd like to be able to denoise pure, random Gaussian noise. We can think of this as starting with a blank canvas $z = \epsilon$ where $\epsilon \sim N(0, I)$ and denoising it to get a clean image $x$.

Repeat the same denoising process as in part 1.2.1, but start with pure noise $\epsilon \sim N(0, I)$ and denoise it for 5 epochs. Display your results after 1 and 5 epochs.

Additionally, compute the average image of the training set. What do you notice between the average image and our attempts to denoise pure noise? Why might this be happening?

We won't provide reference images for this part.

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 ~3 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)
• Sample results on the test set with pure noise $\epsilon \sim N(0, I)$.
• Average image of the training set along with a brief description comparing it to the denoising results.

• 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 Flow Matching Model

For iterative denoising, we need to define how intermediate noisy samples are constructed. The simplest approach would be a linear interpolation between noisy $x_0$ and clean $x_1$ for some $x_1$ in our training data:

Flow can be thought of as the velocity (change in posiiton w.r.t. time) of this vector field, describing how to move from $x_0$ to $x_1$: \begin{equation} u(x_t, t) = \frac{d}{dt} x_t = x_1 - x_0. \tag{B.4}\end{equation}

Our aim is to learn a UNet $u_\theta(x_t,t)$ which approximates this flow $u(x_t, t) = x_1 - x_0$, giving us our learning objective: \begin{equation} L = \mathbb{E}_{x_0 \sim p_0(x_0), x_1 \sim p_1(x_1), t \sim U[0, 1]} \|(x_1-x_0) - u_\theta(x_t, t)\|^2. \tag{B.5} \end{equation}

2.1 Adding Time Conditioning to UNet

Figure 8: Conditioned UNet

Note: It may look like we're predicting the original image in the figure above, but we are not. We're predicting the flow from the noisy $x_0$ to clean $x_1$, which will contain both parts of the original image as well as the noise to remove.

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.

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 $u_\theta(x_t, t)$ to predict the flow at $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: 64.
  • 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-2. 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). Note that if you trained for less than 20 epochs, you will have less steps.
• Sampling results for the time-conditioned UNet for 5 and 10 epochs. We provide up to 20 epochs just for reference.
  • Note: providing a gif is optional and can be done as a bells and whistles below.

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 + t2
# Follow diagram to get the output.
...



        

Training for this section will be the same as time-only, with the only difference being the conditioning vector $c$ and doing unconditional generation periodically.

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 (figure 11). Note that if you trained for less than 20 epochs, you will have less steps.
• Sampling results for the class-conditioned UNet for 5 and 10 epochs. Class-conditioning lets us converge faster, hence why we only train for 10 epochs. Generate 4 instances of each digit as shown above.

Optional Extension: DDPM

This project originally had you implement DDPM instead of flow matching (flow matching can be thought of as a generalization of DDPM). If you'd like to implement DDPM, you can do so by following the original project instructions found here.

Acknowledgements

This project was a joint effort by Ryan Tabrizi, Daniel Geng, and Hang Gao, advised by Liyue Shen, Andrew Owens, Angjoo Kanazawa, and Alexei Efros. We also thank David McAllister and Songwei Yi for their helpful feedback and suggestions.