U-Net & Fully Convolutional Networks

The U-Net Architecture (Ronneberger et al., 2015)

U-Net consists of a Contracting Path (Encoder) to capture context and a symmetric Expanding Path (Decoder) for precise spatial localization:

  CONTRACTING ENCODER (Context)                      EXPANDING DECODER (Localization)
  Input Image (572x572) ──────────────────────────────────────────────────────────► Output Mask (388x388)
     │                                                                                ▲
     ▼                                                                                │
  [ 2x Conv 3x3 ] ─────────────── Skip Connection (Concatenate) ──────────────► [ 2x Conv 3x3 ]
     │ MaxPool 2x2                                                                    ▲ UpConv 2x2
     ▼                                                                                │
  [ 2x Conv 3x3 ] ─────────────── Skip Connection (Concatenate) ──────────────► [ 2x Conv 3x3 ]
     │ MaxPool 2x2                                                                    ▲ UpConv 2x2
     ▼                                                                                │
     └──────────────────────────► [ Bottleneck 2x Conv 3x3 ] ─────────────────────────┘

Encoder-Decoder Mechanics

1. Contracting Encoder: Repeated blocks of $3 \times 3$ convolutions followed by ReLU and $2 \times 2$ Max Pooling downsampling. Feature channels double at each step ($64 \to 128 \to 256 \to 512 \to 1024$).
2. Bottleneck: Deepest layer capturing high-level semantic context.
3. Expanding Decoder: Upsamples feature maps using $2 \times 2$ Transposed Convolutions (Deconvolutions), halving feature channels while doubling spatial dimensions.
4. Skip Connections: Concatenates feature maps from the contracting path directly onto corresponding upsampled decoder feature maps.

Why Skip Connections Are Essential

Pooling operations downsample spatial resolution to capture invariant semantic features.

However, pixel-level segmentation requires exact spatial boundary localization.

Skip connections copy high-resolution edge and boundary details directly from early encoder layers, giving the decoder both what (semantic context) and where (exact spatial location).

FCN: Replacing Fully Connected Layers with 1x1 Convolutions

Traditional classifiers (VGG/AlexNet) end with Dense Fully Connected layers ($4096 \to 1000$).

FCN replaced Dense layers with $1 \times 1$ convolutions:

  Traditional: [ Feature Map 7x7x512 ] ──► Flatten ──► [ Dense Layer 4096 ] (Requires FIXED 224x224 input!)
  FCN:         [ Feature Map 7x7x512 ] ──► [ 1x1 Conv 4096 ]                 (Accepts ANY input H x W!)

Replacing dense layers allows FCNs to take an input image of any resolution and produce a spatially corresponding heat map output in a single forward pass.

Say this out loud

U-Net uses a symmetric U-shaped encoder-decoder architecture with skip connections. The encoder captures deep semantic context through downsampling, while the decoder upsamples feature maps for pixel-wise prediction. Skip connections concatenate early high-resolution spatial maps directly to decoder layers, recovering sharp object boundary details lost during pooling.

Follow-ups to expect

• What is Transposed Convolution (Deconvolution)? A learnable upsampling operation that expands spatial feature dimensions by padding input features with zeros and applying convolutional kernels.
• How is U-Net used in Stable Diffusion? In diffusion models, U-Net takes a noisy image latent vector + timestep embedding + text prompt conditioning (via Cross-Attention) and predicts the exact noise tensor to subtract at each step.