Why U-Net?

Learning Objectives

By the end of this section, you will:

1. Understand why diffusion models need specialized architectures for the image-to-image denoising task
2. Learn why simple CNNs are insufficient for high-quality noise prediction
3. Master the key innovations of U-Net: skip connections and multi-scale processing
4. See how U-Net preserves spatial information through its encoder-decoder structure
5. Connect architecture choices to mathematical requirements of diffusion models

The Big Picture

In the previous chapters, we derived the mathematical framework for diffusion models. We know that training requires a neural network $\epsilon_\theta(x_t, t)$ that:

1. Takes a noisy image $x_t$ and timestep $t$ as input
2. Outputs a prediction of the noise $\epsilon$ that was added
3. Maintains the same spatial resolution as the input (image-to-image mapping)

This seems straightforward: just use a convolutional neural network, right? Not quite. The denoising task has unique requirements that demand a specialized architecture.

The U-Net architecture, originally designed for medical image segmentation in 2015, turned out to be perfectly suited for diffusion models. Its encoder-decoder structure with skip connections provides exactly what we need: multi-scale feature extraction combined with precise spatial localization.

The Image-to-Image Problem

Diffusion models perform dense prediction: for every input pixel, we must predict the corresponding noise value. This is fundamentally different from classification (one label per image) or detection (bounding boxes).

Noise prediction shares characteristics with segmentation (per-pixel output) but has additional requirements:

• Precise spatial alignment: The predicted noise must align exactly with the input image. Off-by-one errors create artifacts.
• Multi-scale understanding: Low-frequency noise patterns span large regions; high-frequency noise is pixel-level.
• Continuous values: Unlike classification (discrete), noise prediction is regression. Small errors accumulate across sampling steps.

Why Not a Simple CNN?

Let's consider what happens if we try to use a simple stack of convolutional layers:

1import torch
2import torch.nn as nn
3
4class SimpleNoisePredictor(nn.Module):
5    """A naive approach that won't work well for diffusion."""
6
7    def __init__(self, in_channels: int = 3, hidden_dim: int = 64):
8        super().__init__()
9
10        # Just stack convolutions - no skip connections, no multi-scale
11        self.layers = nn.Sequential(
12            nn.Conv2d(in_channels, hidden_dim, 3, padding=1),
13            nn.ReLU(),
14            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
15            nn.ReLU(),
16            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
17            nn.ReLU(),
18            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
19            nn.ReLU(),
20            nn.Conv2d(hidden_dim, in_channels, 3, padding=1),  # Output same size
21        )
22
23    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
24        # Problem: t is not used! How does the network know the noise level?
25        return self.layers(x)

This simple approach fails for several reasons:

Problem 1: Limited Receptive Field

Each 3x3 convolution only sees a 3x3 neighborhood. To understand global context (what objects are in the scene), we need much larger receptive fields. With 5 layers of 3x3 convolutions, the receptive field is only 11x11 pixels—far too small for 256x256 or 512x512 images.

Problem 2: No Multi-Scale Processing

Noise exists at multiple scales: large smooth regions of noise, medium texture-like patterns, and fine pixel-level variations. A fixed-resolution network cannot efficiently process all scales simultaneously.

Problem 3: Information Loss

As information flows through many convolutional layers, fine spatial details get averaged out. By the time we reach the output, we've lost the precise edge locations needed for sharp reconstructions.

Problem 4: No Time Conditioning

The network doesn't know what timestep $t$ it's denoising. But the noise statistics change dramatically from $t=1000$(pure noise) to $t=1$ (almost clean image). The network needs to adapt its behavior based on $t$.

U-Net's Key Innovations

The U-Net architecture, proposed by Ronneberger et al. in 2015 for biomedical image segmentation, introduced two key innovations that make it perfect for diffusion:

1# Conceptual U-Net structure for diffusion
2
3class DiffusionUNet:
4    """
5    U-Net for diffusion models combines:
6    1. Encoder (downsampling): Extract multi-scale features
7    2. Bottleneck: Global context at lowest resolution
8    3. Decoder (upsampling): Reconstruct with skip connections
9    4. Time conditioning: Adapt to noise level
10    """
11
12    # ENCODER: Progressively downsample
13    # Input: [B, 3, 256, 256]
14    # -> Down1: [B, 64, 128, 128]   (extract edges, textures)
15    # -> Down2: [B, 128, 64, 64]    (extract patterns)
16    # -> Down3: [B, 256, 32, 32]    (extract objects)
17    # -> Down4: [B, 512, 16, 16]    (extract scenes)
18
19    # BOTTLENECK: Process at lowest resolution
20    # [B, 512, 16, 16] -> [B, 512, 16, 16]
21    # Global context, self-attention possible here
22
23    # DECODER: Progressively upsample with skip connections
24    # Up1: concat(Down4) -> [B, 256, 32, 32]
25    # Up2: concat(Down3) -> [B, 128, 64, 64]
26    # Up3: concat(Down2) -> [B, 64, 128, 128]
27    # Up4: concat(Down1) -> [B, 64, 256, 256]
28
29    # OUTPUT: Project to noise prediction
30    # [B, 64, 256, 256] -> [B, 3, 256, 256]
31
32    # TIME CONDITIONING: Injected at every block
33    # t -> sinusoidal embedding -> MLP -> add/scale features

The architecture gets its name from the U-shape when visualized: the encoder goes down the left side, the bottleneck is at the bottom, and the decoder goes up the right side.

Skip Connections

Skip connections are the defining feature of U-Net. They directly connect encoder layers to corresponding decoder layers at the same resolution:

• Preserve spatial information: High-resolution features from the encoder (edges, textures) are passed directly to the decoder, bypassing the bottleneck.
• Enable gradient flow: Gradients can flow directly from decoder to encoder through skip connections, enabling training of very deep networks.
• Combine semantics with details: The decoder combines semantic features from the upsampling path with spatial details from skip connections.

Mathematically, if $f_{enc}$ are encoder features and$f_{dec}$ are upsampled decoder features, the combined features are:

$f_{combined} = \text{Concat}(f_{enc}, f_{dec})$

This doubles the channel count, which is then reduced by a convolution in the decoder block.

Multi-Scale Feature Processing

The encoder-decoder structure enables multi-scale processing:

Each resolution level has a different receptive field relative to the original image. At 16x16, each spatial location corresponds to a 16x16 patch in the original image, enabling global reasoning. At 256x256, each location sees only a small neighborhood, enabling precise local predictions.

Architecture Overview

Let's visualize the complete U-Net architecture for diffusion models. Click on any block to learn more about its function:

The key components we'll implement in the following sections:

1. ResBlocks: The building blocks with residual connections
2. Downsampling: Strided convolutions or pooling to reduce resolution
3. Upsampling: Transposed convolutions or interpolation to increase resolution
4. Time Embedding: Sinusoidal encoding of timestep injected into each block
5. Attention Layers: Self-attention for global context (especially at low resolutions)
6. Skip Connections: Concatenation of encoder features to decoder

Summary

In this section, we learned why the U-Net architecture is the ideal choice for diffusion models:

1. Image-to-image nature: Noise prediction requires the same spatial resolution as input, with per-pixel accuracy.
2. Simple CNNs are insufficient: They lack multi-scale processing, have limited receptive fields, and lose spatial details.
3. Skip connections preserve details: High-resolution features from the encoder flow directly to the decoder, maintaining spatial precision.
4. Multi-scale processing: The encoder-decoder structure naturally processes noise at multiple frequency scales.
5. Gradient flow: Skip connections enable training very deep networks by providing gradient highways.