U-Net

Chapter 1: The Encoder–Decoder — Down, Then Up

The first half of U-Net's idea is the encoder–decoder shape. It has two paths. The encoder (or contracting path) progressively shrinks the image — convolve, then downsample, again and again — squeezing a big detailed picture into a small stack of abstract features. The decoder (or expanding path) does the reverse — upsample, then convolve — growing those abstract features back into a full-resolution output. Down to understand, up to localize.

Why go down first

Each downsampling step does two things at once. It shrinks the spatial size (so the next layers are cheaper), and crucially, it widens the receptive field — how much of the original image each feature “sees.” After a few downsamples, a single feature summarizes a large region of the image, so it can encode high-level meaning: “this region is a cell nucleus.” Without going down, features stay local and the network never grasps context. The encoder trades resolution for understanding.

Why go up after

The decoder reverses the journey. It takes the small, semantically-rich feature map at the bottom and progressively upsamples it — doubling the spatial size each step, applying convolutions to refine — until it's back to the original image resolution. Now the output is the same size as the input, so it can carry a prediction for every pixel. The encoder built up meaning by shrinking; the decoder rebuilds spatial precision by growing. Together they form the two strokes of the letter U — down the left side, up the right.

The data flow, traced

See it: the resolution pyramid

The widget shows feature-map sizes through a plain encoder–decoder. Step through: watch the spatial size halve on the way down (and channels grow), reach the tiny bottleneck, then double back up to full resolution. This down-then-up profile is the skeleton of U-Net — but as the next chapter shows, the skeleton alone produces a disappointingly blurry result.

Chapter 2: The Blur — Why Encoder–Decoder Isn’t Enough

Here's the problem that makes plain encoder–decoders fail, and it sets up U-Net's one big idea. The decoder has to reconstruct a full-resolution, pixel-precise output — but the only thing it receives is the tiny bottleneck feature map. And that bottleneck threw away almost all the fine spatial detail on the way down. The result: a blurry, smeared output that gets the rough shape right but the exact boundaries wrong.

Why the bottleneck can't carry detail

Think about the numbers. A 256×256 image has 65,000 pixels of spatial detail. The bottleneck might be 16×16 — just 256 spatial locations. Even though each location is rich in channels (meaning), there are simply not enough spatial positions to encode where every fine edge was. The encoder deliberately discarded that — pooling a 2×2 region into one value forgets which of the four pixels was the bright one. So when the decoder upsamples, it has to guess the fine detail, and its best guess is a smooth, blurry average. The sharp boundary of a cell becomes a fuzzy gradient.

The tension, made sharp

So we're stuck between two needs that pull in opposite directions. The encoder must downsample to build context (what the object is). But downsampling destroys the fine detail the decoder needs for precise boundaries. A plain encoder–decoder gives you context at the cost of precision — blurry masks. We need a way to give the decoder both: the high-level meaning from the deep bottleneck and the fine detail from the shallow, high-resolution early layers. Hold those two streams of information until they can be combined.

See it: reconstruction from the bottleneck alone

The widget shows a crisp input mask, then what a plain encoder–decoder reconstructs from the bottleneck alone. Crank the bottleneck size down (more downsampling) and watch the output get blurrier — the boundaries smear because the spatial detail was pooled away. This blur is exactly the disease skip connections cure.

Chapter 3: Skip Connections — The One Big Idea

This is the heart of U-Net, the idea that turns a blurry encoder–decoder into a boundary-sharp segmenter. The insight: the fine detail the decoder needs still exists — in the encoder's early layers, before downsampling threw it away. So instead of forcing all information through the lossy bottleneck, copy each encoder feature map directly across to the decoder at the same resolution. These shortcuts are the skip connections, and they're what put the cross-bars in the letter U.

How a skip connection works

At each level of the decoder, just before refining, you have two feature maps at the same spatial size: the upsampled one coming up from below (rich in context, but coarse), and the copied one skipping across from the matching encoder level (rich in fine detail, high-resolution). U-Net concatenates them — stacks them along the channel dimension — then runs convolutions over the combined stack. Now the decoder's refinement step can use both the high-level meaning and the precise spatial detail at once. Context from below, detail from the side, fused.

Worked example: concatenation

At decoder level 2, suppose the upsampled feature map from below is 128×128 with 256 channels (carrying context), and the encoder's level-2 feature map — copied across the skip — is also 128×128 but with 128 channels (carrying fine detail). Concatenating along channels gives a 128×128 map with 256 + 128 = 384 channels. The next convolution sees all 384 channels together, so a single output pixel can be computed from both the “this region is a cell” signal and the “the edge is exactly here” signal. The spatial dimensions must match for concatenation to work — which is exactly why the encoder and decoder are mirror images with matching resolutions at each level.

From scratch: the skip in code

python
def unet_forward(x):
    # encoder — SAVE each feature map before downsampling
    e1 = enc_block1(x);        s1 = e1;  e1 = down(e1)   # save s1 (high-res)
    e2 = enc_block2(e1);       s2 = e2;  e2 = down(e2)   # save s2
    b  = bottleneck(e2)                                  # deepest, lowest-res
    # decoder — CONCATENATE the saved skip at each level
    d2 = up(b)
    d2 = dec_block2(torch.cat([d2, s2], dim=1))   # ← skip: context + detail
    d1 = up(d2)
    d1 = dec_block1(torch.cat([d1, s1], dim=1))   # ← skip again
    return final_conv(d1)                            # full-res mask

The whole U-Net difference is those two torch.cat([..., s], dim=1) calls. The encoder saves its feature maps (s1, s2) before downsampling; the decoder concatenates them back in at the matching resolution. Without those lines you have a plain (blurry) encoder–decoder. With them, you have U-Net.

See it: skip on vs. off

The same image, segmented with skip connections on and off. With skips off, the output is the blur from Chapter 2 — context but no precise boundaries. Toggle skips on and the high-resolution detail floods back in: crisp, accurate edges. This single toggle is the difference between a useless segmenter and a state-of-the-art one.

Chapter 4: The Full U — Watch It Flow

Now see the whole architecture in motion. This is a complete U-Net: an image enters at the top-left, flows down the contracting encoder (blocks shrinking, getting deeper), across the bottleneck at the bottom, then up the expanding decoder (blocks growing back), and out as a segmentation mask at the top-right. The horizontal arrows across the middle are the skip connections, carrying each encoder level's detail straight to the matching decoder level. That shape — down, across, up, with rungs between — is why it's called U-Net.

Press Run to send data through and watch the flow. Then toggle the skip connections and re-run: with skips on, the output mask is crisp; with skips off, the same network produces the familiar blur. You can see the detail traveling across the rungs of the U.

No quiz — the flow is the test. If you can trace why turning off the rungs blurs the output, you understand U-Net's core.

Chapter 5: Upsampling — How the Decoder Grows

We've talked about the decoder “upsampling” — growing a small feature map back to a larger one. But how do you actually add spatial resolution? You're conjuring more pixels than you started with, so something must fill in the gaps. There are two main approaches, and the choice has real consequences for the output's quality.

Method 1: interpolation + convolution

The simple approach: blow the feature map up with plain image resizing — nearest-neighbor (copy each value into a 2×2 block) or bilinear (smoothly interpolate between values) — then run a normal convolution to refine. The resize adds the resolution (dumbly), and the convolution learns to clean it up. It's simple, robust, and the modern default in many architectures.

Method 2: transposed convolution

The learned approach: a transposed convolution (sometimes loosely called “deconvolution”) does the upsampling and the learning in one step. Instead of sliding a filter to shrink, it does the reverse — each input value is multiplied by a learned filter and painted into a larger output region, with overlapping regions summed. The network learns how to expand. More flexible, but with a notorious failure mode.

Worked example: why the checkerboard appears

Imagine a transposed convolution with a filter of size 3 and a stride of 2 (a common setting that causes trouble). As the filter steps across the output by 2 but spans 3, adjacent output positions receive overlapping contributions — but unevenly. Output pixel A might be covered by 2 filter placements while neighbor B is covered by only 1. So A accumulates roughly twice the signal of B, every other pixel, in both directions — a checkerboard. The fix is to make the filter size a multiple of the stride (so coverage is even), or to avoid transposed convs entirely and use resize-then-convolve, where every output pixel is treated identically.

See it: three ways to upsample

A small feature map grown four ways: nearest-neighbor (blocky), bilinear (smooth), transposed-convolution done well (even coverage), and transposed-convolution done badly (the checkerboard). Toggle between them to see the tradeoffs — and to recognize the checkerboard artifact when you spot it in a real generated image.

Chapter 6: Training a U-Net — The Loss Matters

A U-Net outputs a per-pixel prediction, so the obvious loss is per-pixel cross-entropy: classify each pixel (tumor or not) and average the loss over all pixels. It works — but in segmentation there's a vicious trap that makes this loss alone dangerous: class imbalance.

The imbalance trap

In medical images, the thing you care about is often tiny. A tumor might be 2% of the pixels; the other 98% are healthy background. Now watch what per-pixel accuracy does: a model that predicts “background” for every single pixel — finding nothing at all — scores 98% accuracy. The loss is happily low. The model has learned to do nothing, and the metric congratulates it. Cross-entropy, averaged over pixels, barely notices the tumor because it's drowned out by the sea of easy background pixels.

The fix: Dice loss

The standard answer is the Dice coefficient (and its loss form). Dice measures overlap: roughly, twice the area where prediction and ground-truth tumor agree, divided by the total tumor area in both. It's 1 when the predicted mask perfectly matches the true mask, and 0 when they don't overlap at all — and crucially, it ignores the background entirely. Predict all-background and Dice is 0 (no overlap with the tumor), not 0.98. The model can no longer cheat by ignoring the rare class. In practice people often combine cross-entropy (stable gradients everywhere) with Dice (focus on the overlap) to get the best of both.

Worked example: the two losses disagree

A 100-pixel image with a 4-pixel tumor. The model predicts all-background (finds nothing). Pixel accuracy: 96 of 100 pixels correct = 96% — looks great. Dice: the overlap between the predicted tumor (empty) and the true tumor (4 pixels) is zero, so Dice = 2×0 / (0 + 4) = 0 — correctly screams “you found nothing.” Same prediction, and the two scores couldn't disagree more. Train on Dice (or Dice + CE) and the model is forced to actually locate the tumor, because all-background scores zero.

See it: accuracy vs. Dice as the tumor shrinks

A model that always predicts “background.” Shrink the true tumor and watch pixel accuracy soar toward 100% (it's right about all the background) while Dice stays pinned at 0 (it never overlaps the tumor). The gap between the two is the imbalance trap made visible — and the reason segmentation is trained on overlap losses, not accuracy.

Chapter 7: U-Net as the Diffusion Backbone

Here's why U-Net matters far beyond medical imaging. When diffusion models — Stable Diffusion, the original DALL·E 2, Imagen — took over image generation, the network at their core, the one doing the actual work, was almost always a U-Net. The same shape you just learned for segmenting cells became the engine of AI art. Understanding why reveals something deep about what U-Net is really good at.

What diffusion asks the network to do

A diffusion model generates an image by starting from pure noise and removing a little noise at a time, over many steps, until a clean image emerges. The network's job at each step is: given a noisy image, predict the noise that was added (so it can be subtracted). Look at the shape of that task — the input is a full-resolution image, and the output is also a full-resolution image (the predicted noise). Image in, image out, same size. That is exactly the shape U-Net was built for. Segmentation maps pixels to pixel-labels; denoising maps pixels to pixel-noise. Same architecture, different output meaning.

Two additions: time and text

Diffusion U-Nets add two things to the segmentation U-Net. First, time conditioning: the network must know which denoising step it's on (early steps remove coarse noise, late steps refine fine detail), so the timestep is encoded and injected into every block. Second, cross-attention to text: for text-to-image, attention layers are inserted so the denoising can be guided by a prompt (“a cat in a hat”). The U-shape with skips stays exactly the same — these are conditioning signals threaded through the familiar backbone.

The bridge to DiT

Recently, some models replaced the U-Net backbone with a pure transformer — the Diffusion Transformer, or DiT, which powers Sora and newer image models. DiT chops the image into patches and runs attention, scaling more smoothly than convolutional U-Nets. But the conceptual job is identical (noise in, noise out), and many production systems still use U-Nets or U-Net/transformer hybrids. U-Net was the backbone that made the diffusion era possible; DiT is its successor, but the lineage is direct. (See the DiT and Diffusion lessons.)

See it: denoising with a U-Net

Step through a diffusion process. Start from pure noise; at each step the U-Net predicts the noise, it's subtracted, and a clearer image emerges — until a clean result appears. Watch the same “image in, image out” network you learned for segmentation, now generating.

Chapter 8: The U-Net Family

U-Net's skeleton — contract, bottleneck, expand, with skip connections — turned out to be so good that it spawned a whole family of variants, each adapting the template to a new domain or fixing a limitation. Knowing them shows just how general the core idea is.

The major variants

• 3D U-Net / V-Net: swap 2D convolutions for 3D ones to segment volumes — CT and MRI scans are 3D stacks of slices. The U-shape is identical; the operations just work in three dimensions. V-Net added residual connections and popularized the Dice loss.
• Attention U-Net: add attention gates on the skip connections, so the decoder learns to focus the incoming detail on relevant regions and suppress irrelevant background. The skips become selective rather than copying everything.
• nnU-Net: not a new architecture but a brilliant auto-configuring framework — it inspects your dataset and automatically picks the U-Net depth, patch size, and training recipe. It famously beats fancier custom models on medical benchmarks, proving a well-tuned plain U-Net is remarkably hard to beat.
• Transformer / ConvNeXt U-Nets: replace the convolutional blocks with transformer blocks (TransUNet, Swin-UNet) or modernized convolutions, keeping the U-shape and skips but upgrading what fills them. The diffusion U-Nets from Chapter 7 are this kind of hybrid — convolutions plus attention.

See it: what each variant changes

Select a variant and see which part of the base U-Net it modifies (highlighted), while everything else stays the same. It's a vivid reminder that these are all the same architecture with one targeted change — the skeleton is shared, the modification is local.

Chapter 9: Connections & Cheat Sheet

You now understand U-Net from the inside: why pixel-level prediction needs both context and precision, how the encoder–decoder captures context, why the bottleneck alone produces blur, how skip connections rescue the fine detail, how upsampling works, how to train it with overlap losses, why it became the diffusion backbone, and how its variants all share one skeleton. The thread: map a high-resolution input to a high-resolution output by going down for meaning and up for precision — and carry the detail across so precision survives the round trip.

The cheat sheet

A decision guide

Where this connects

• Skip Connections — the residual cousin of U-Net's skips; both route information around bottlenecks (additive vs concatenative).
• Diffusion Models — U-Net is the engine that predicts noise at each denoising step.
• DiT (Diffusion Transformer) — the transformer successor that's replacing U-Net backbones in newer generative models.
• Detection & Segmentation — U-Net's home turf, alongside other dense-prediction architectures.
• Data Augmentation — elastic deformations were key to U-Net working on tiny medical datasets.
• Loss Functions — Dice, IoU, and cross-entropy for imbalanced dense prediction.
• Vision Transformers — the backbone now fused into transformer U-Nets.

“To see clearly, you must first step back to take in the whole — then return to attend to every detail.” That round trip, with memory carried across, is the U.