U-Net, Section by Section

The shape problem

Many vision tasks need a per-pixel output, not a single class per image: medical image segmentation, satellite land-cover parsing, depth estimation, and most recently denoising in diffusion models. The output has the same spatial size as the input.

A pure CNN classifier squashes spatial resolution down to a tiny feature map (good for "is this a cat?"). To produce a full-resolution mask you have to get the resolution back — and the only honest way to do that is to have remembered it on the way down.

The architecture, drawn to scale

A 4-level U-Net: 4 encoder blocks (each halves spatial size and doubles channels), one bottleneck, 4 decoder blocks (each doubles spatial size, halves channels, and concatenates the matching encoder feature). For a 64×64 input the bottleneck is 4×4. That's 16× spatial compression in two dimensions, or 256× total — enormous information loss without skips.

Forward pass on a real synthetic image

We synthesise an input: two coloured ellipses on noise. The target mask labels which pixels belong to ellipse 1 vs ellipse 2 vs background. A small U-Net (random weights; we hand-craft the conv kernels for visibility) processes the image. Inspect the feature map at each scale:

Toggle the skips, watch detail die

Each skip connection at level $\ell$ concatenates the encoder's level-$\ell$ feature map onto the decoder's input at the same level. Turn one off and the corresponding scale of detail goes missing. Turn them all off and you've built a plain encoder-decoder, which is exactly the architecture U-Net was invented to beat.

Why concatenate, not add?

ResNet adds: $y = F(x) + x$. U-Net concatenates: $y = [F(x); \text{skip}]$. The reason is information preservation. The encoder feature carries spatial detail the bottleneck has discarded. Adding a vector with very different statistics into the bottleneck output would interfere with the decoder's learning. Concatenation lets the decoder freely choose which channels to read.

This costs more channels in the decoder convolutions, but for dense prediction tasks it's worth it. Modern U-Nets (and diffusion U-Nets, see Step 5) sometimes mix — add for residuals within a block, concatenate for skips across the U.

The U inside diffusion

Stable Diffusion, Imagen, every modern image-diffusion backbone is a U-Net. The denoiser takes a noisy image and predicts the noise it has to remove. That's a per-pixel regression problem — same shape, same downsampling and upsampling story, plus three additions:

• Time-step conditioning. The current noise level $t$ is fed into every block via a sinusoidal embedding (see the positional encoding section in the attention article for the recipe).
• Cross-attention to text. Inside each decoder block the U-Net does cross-attention into a CLIP-encoded prompt, so generation is text-conditional.
• Self-attention at the bottleneck. After spatial resolution is small (8×8 ish), self-attention connects far-apart features — locality bias disappears.

The skips do the same job they always did: they preserve high-frequency detail that the encoder has thrown away. That's why diffusion outputs have crisp edges instead of looking like JPEG-compressed dreams.

The variants worth knowing

• U-Net (original). Ronneberger et al., 2015. Built for biomedical segmentation; trained on a few hundred labelled cells. Still the baseline 10 years later.
• nnU-Net. Self-configuring U-Net; specifies preprocessing, architecture, and training schedule from the dataset. Auto-tunes for your modality. The thing to use first on any segmentation task.
• 3-D U-Net. Same architecture with 3-D convolutions; volumetric medical (CT, MRI), satellite time-cubes, point-cloud voxelisations.
• Attention U-Net. Gates each skip connection by a learned attention mask. Helps when the image has many small targets.
• TransUNet / Swin-UNet. Swap the encoder for a Vision Transformer; keeps the U for resolution recovery. The standard recipe for medical image segmentation in 2024+.
• U-Net++ (NestedU-Net). Extra dense skips at every intermediate scale. Slightly stronger, slower.
• Diffusion U-Net (DDPM, Stable Diffusion). U-Net + time-step embedding + cross-attention to text. Same skips as in 2015.

Diffusion's U-Net — why an old architecture won the generative race

Every Stable Diffusion checkpoint is a U-Net, lightly modernised: the encoder + bottleneck + decoder structure is the original 2015 design, with three additions:

• Time-step embedding. A sinusoidal time embedding $t$ is mapped to a per-layer scale + bias (FiLM-style). The same network does denoising at all noise levels by reading $t$.
• Cross-attention to a text condition. At a few mid-resolution layers, the spatial feature map attends to a tokenised text prompt's embeddings. This is where "prompt steering" actually happens.
• Self-attention blocks at low resolution. The U's bottleneck has full self-attention over the spatial grid (cheap because grid is small there); high-resolution layers stay purely convolutional.

In 2024–2026 the field has moved partially toward DiT (Diffusion Transformers) — pure transformer backbones with no U-Net skips. DiTs scale better past ~2B parameters; U-Nets still dominate sub-1B image-generation models because they're much more parameter-efficient on small data.

Practical implementation notes

• Input size must be divisible by $2^D$ where $D$ is the depth. A 4-stage U-Net needs inputs of $H, W$ divisible by 16. Pad before, crop after.
• Padding choices. Original used "valid" (no padding); modern implementations use "same" with reflection or zero padding to keep spatial dimensions stable through skips. Reflection padding usually gives slightly cleaner edges.
• Upsampling: ConvTranspose vs interpolate-then-conv. ConvTranspose can produce checkerboard artefacts at certain stride / kernel ratios. Most modern U-Nets use bilinear / nearest-neighbour upsampling followed by a 3×3 conv — visually cleaner.
• Number of stages. 4–5 stages cover most natural-image tasks. Deeper U's need wider channels (the bottleneck width controls capacity).
• Loss. Cross-entropy + Dice for medical / class-imbalanced segmentation; MSE for diffusion denoising; perceptual losses (VGG features) for image-to-image translation.

Reading list

• Ronneberger, Fischer, Brox (2015) — U-Net. The original.
• Çiçek et al. (2016) — 3D U-Net.
• Oktay et al. (2018) — Attention U-Net.
• Isensee, Jaeger, Kohl, Petersen, Maier-Hein (2021) — nnU-Net. Self-configuring U-Net for medical segmentation.
• Ho, Jain, Abbeel (2020) — Denoising Diffusion Probabilistic Models. The U-Net inside diffusion.
• Peebles & Xie (2022) — Scalable Diffusion Models with Transformers (DiT). The eventual transformer-based successor.