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) reverses the change in spatial scale — upsample, then convolve — growing those abstract features back into a dense output. Upsampling is not an inverse of the encoder.

Why go down first

Downsampling reduces the number of spatial locations processed by later layers. It also increases the spacing, measured in input pixels, between neighboring features. Subsequent convolutions therefore gather context over a larger input region. For example, a 3×3 convolution increases the receptive-field width by 2 pixels at stride 1, but by 8 input pixels after two factor-of-two reductions. This is efficient receptive-field growth, not a guarantee that learned features understand an object.

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 reaches the desired output resolution. Now the output can carry a prediction at each desired location; same-padding variants can match the input size. 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

For a stride-1 3×3 convolution, padding 1 preserves size; padding 0 removes one pixel at each border. The Conv2d shape formula makes this explicit.

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 — the next chapter examines why a narrow deep path can make boundary reconstruction harder.

Chapter 2: The Blur — What a Bottleneck Can Lose

A decoder without long skips receives its information through the deepest feature map. If the encoder has discarded boundary details, enlarging that map cannot recover those details uniquely. Predictions may then miss thin structures or localize boundaries poorly. This is a motivation for U-Net skips, not a theorem that every plain encoder–decoder must blur.

What pooling actually loses

Take two 2×2 patches: one has a bright top-left pixel; the other has a bright bottom-right pixel. Both have maximum 1 and average 0.25. Max pooling or average pooling on that single channel cannot distinguish them. These operations are non-injective: different inputs can produce the same output. A decoder given only that pooled value cannot know which patch it came from.

The tension, made sharp

The encoder develops a broad view, while the decoder needs local evidence. Long skips provide a direct route for earlier, higher-resolution features, reducing reliance on the deepest representation. Those features are learned responses, not a separate set of perfect edge labels. A trained decoder must learn how to use both streams.

See it: reconstruction from the bottleneck alone

The widget performs a specific, deliberately simple operation: average a 24×24 synthetic mask into a smaller grid, then bilinearly resize it. Smaller grids retain less of this mask’s detail. This demonstrates information loss and interpolation, not the output of a trained encoder–decoder.

Chapter 3: Skip Connections — The One Big Idea

U-Net gives the decoder a second source of information: feature maps saved from the encoder before each downsampling stage. It copies those maps across to the corresponding decoder stage and fuses them with upsampled deeper features. These long skip connections are the cross-bars of the U. They help localization by making local evidence accessible; they do not guarantee a perfect boundary.

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 — so implementations explicitly align the grids by design, cropping, or padding.

From scratch: the skip in code

python
# Schematic: NCHW tensors; same-padding blocks and matching spatial sizes.
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)                            # per-pixel logits; apply sigmoid/softmax afterward

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 encoder–decoder with no long skips. With them, you have U-Net.

See it: skip on vs. off

This toy comparison starts with the same synthetic mask as Chapter 2. The no-skip view shows pooling plus interpolation. The skip-on view directly copies the original mask to illustrate access to a high-resolution route. A real U-Net concatenates learned features and learns a decoder; it does not copy the ground-truth mask or become accurate just because a switch is enabled.

Chapter 4: The Full U — Watch It Flow

Trace an illustrative same-padding U-Net: input at top-left, contracting encoder, bottleneck, expanding decoder, and prediction at top-right. The horizontal rungs carry encoder features to matching decoder scales. The two bottom boxes depict the entrance and exit of one bottleneck stage, not two separately skipped stages. Channel counts and exact operator choices remain implementation decisions.

Press Run to trace the deep path, then toggle the skips. The output thumbnails are schematic indicators of detail access, not network predictions or evidence of an accuracy improvement. The architectural question is whether each decoder stage can use both its deeper input and the corresponding encoder features.

No quiz — trace both routes into a decoder block, explain what concatenation changes, and distinguish this same-padding diagram from the original cropped-skip design.

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 refine the resized features. It is a common alternative to learned strided upsampling.

Method 2: transposed convolution

A transposed convolution applies the transpose of the matrix describing a convolutional linear map. Operationally, each input value contributes a learned kernel to the output; overlapping contributions are summed. With stride greater than 1 it can increase spatial size. It is sometimes called “deconvolution,” but it is not generally an inverse that reconstructs the original signal. See the PyTorch operator definition.

Worked example: why the checkerboard appears

For a one-dimensional all-ones input and all-ones kernel of size 3, stride 2, interior overlap counts alternate between 1 and 2. In two dimensions those counts multiply: interior positions receive 1, 2, or 4 contributions. This is a concrete arithmetic example, not a universal brightness rule for learned filters. Size 4, stride 2 removes this interior count imbalance; output boundaries still have fewer contributions.

See it: interpolation and overlap

The first two modes resize a 4×4 feature map to 16×16 using nearest-neighbor or bilinear interpolation. The other two compute actual all-ones transposed-convolution overlap counts at stride 2, with kernel size 4 or 3. They isolate coverage, not learned image quality. Border effects remain visible in both count maps.

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

A foreground region might occupy 2% of the pixels. Predicting background everywhere then gives 98% pixel accuracy while missing the whole region. This does not imply low cross-entropy: CE uses probabilities, and a confident false negative has a large penalty. Averaging over many background pixels can make class balance matter, but accuracy and CE are different quantities.

Measuring overlap: Dice

For a binary foreground mask, Dice is twice the intersection divided by the sum of predicted and true foreground sizes: 2TP / (2TP + FP + FN). It ignores true negatives, but false-positive foreground predictions increase the denominator and reduce the score. When the true foreground is nonempty, an all-background prediction scores 0. If both masks are empty, the ratio is 0/0; implementations must specify an empty-mask convention or smoothing.

Worked example: scores versus losses

For the 100-pixel example, an empty predicted foreground has Dice 0 and pixel accuracy 96%. If the model assigns foreground probability 0.01 everywhere, mean binary CE is −[0.04 log(0.01) + 0.96 log(0.99)] ≈ 0.194. Decreasing that probability toward zero makes the missed foreground more confidently wrong and increases its penalty. For training, a soft Dice loss replaces hard masks with probabilities; its smoothing and reduction across images/classes must be chosen explicitly. Combining CE and soft Dice is one common option, not a guarantee of good segmentation.

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 accuracy should not be the only evaluation metric. Accuracy is also not a differentiable training objective.

Chapter 7: U-Net as the Diffusion Backbone

U-Net became an important backbone beyond segmentation. The DDPM paper used a U-Net-like network to model denoising, while latent diffusion adapted this idea to compressed image representations. The transferable feature is multiscale, spatially aligned prediction. Diffusion is a process; U-Net is one possible neural architecture inside it.

What diffusion asks the network to do

In the noise-prediction parameterization, training constructs x_t = √ᾱ_t x₀ + √(1 − ᾱ_t) ε, where ε is sampled Gaussian noise and ᾱ_t sets the signal level. The network estimates ε from x_t, timestep t, and any conditioning. The sampler then uses that prediction with schedule-dependent coefficients, and sometimes fresh noise, to produce the next state. It does not simply subtract the raw prediction. Other parameterizations predict the clean sample x₀ or a velocity-like target v.

Two additions: time and text

Time conditioning tells the denoiser the current noise level, often through an embedding injected into residual blocks. A text-conditioned latent diffusion model can also use cross-attention: spatial features query an encoded prompt. Conditioning can instead be class labels, another image, or be absent. Exact block layouts vary; time and text do not make every diffusion architecture an unchanged segmentation U-Net.

The bridge to DiT

The DiT paper explored transformers over patches of a latent representation as an alternative diffusion backbone. Its experiments studied scaling model size and compute. This shows that denoising does not require a convolutional U-shaped hierarchy. U-Nets, transformers, and hybrids are design choices; no architecture is a universal successor or winner for every task.

See it: denoising with a U-Net

The widget is a visual analogy for noise giving way to structure. It linearly blends a fixed random grid with a known synthetic target over eight steps. It runs no denoising network, has no diffusion noise schedule, and does not demonstrate generation quality. Use it to picture repeated refinement while keeping the actual training target and sampler distinct.

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. These are distinct volumetric architectures, not identical implementations. 3D U-Net adapts the encoder–decoder to 3D; V-Net uses volumetric residual blocks and a Dice-based objective.
• 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: an auto-configuring segmentation framework — it inspects your dataset and automatically picks the U-Net depth, patch size, and training recipe. The nnU-Net paper reports strong results across medical segmentation tasks. That supports careful configuration as a baseline, not a promise to win on every dataset.
• Transformer hybrids: TransUNet combines a CNN feature extractor with a transformer encoder and an upsampling decoder that uses skip features. It does not replace every convolutional block with attention. Other U-shaped transformer models use different encoder, decoder, and skip designs. Diffusion U-Nets can also combine convolutions and attention.

See it: what each variant changes

Select a variant to highlight its main architectural emphasis. The transformer view marks a hybrid encoder stage rather than claiming that all blocks become transformers. These are schematic comparisons; consult each paper for its full implementation.

Chapter 9: Connections & Cheat Sheet

U-Net makes a practical bet: dense prediction benefits from combining features computed at several spatial scales. The deep route gathers context; long skips make earlier features available to the decoder. That is a useful template for segmentation and denoising, with performance determined by data, representation, optimization, and implementation details as well as the diagram.

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 — the process and objectives that a U-Net or another denoiser can implement.
• DiT (Diffusion Transformer) — a transformer alternative to a U-Net denoising backbone.
• 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.