U-Net: Convolutional Networks for Biomedical Image Segmentation

Pooling forgets where things are; the skip connections remember.

A classifier crushes an image down to a single label. Segmentation needs the opposite: a label for every pixel. U-Net does both by shrinking an image to read it, then growing it back to full resolution, with shortcuts that hand the lost detail straight across.

How do you train a network to label every pixel of a microscope image when you only have thirty labelled images to learn from?

For most of deep learning's short history, a convolutional network answered one question per image: what is this? You feed in a picture and out comes a single label. Cat. The whole image collapses to one word. That is image classification, and by 2015 convnets were very good at it.

A biologist staring at an electron-microscope image of brain tissue needs a different kind of answer. Not "this image contains cells," but "this pixel is the inside of a cell, that pixel is a membrane, and the hairline between these two cells that are pressed together is the wall that separates them." A class label for every single pixel. That task is called semantic segmentation, and it is harder than it looks for a structural reason.

To recognize what something is, a network needs context. It stacks convolutions and pooling layers, and each $2\times2$ pooling step replaces a small patch of the feature map with a single value. A neuron deeper in the stack therefore sees a wider region of the image and builds a more abstract feature, but it can no longer say which exact pixel that feature came from. That is the whole trick of an image classifier, and it is the wrong thing to do if you also need to know where a boundary sits down to the pixel. You gain what at the cost of where. Context and localization pull against each other, and this is the tension the paper is built to resolve.

U-Net, out of Olaf Ronneberger's group at the University of Freiburg, resolves it with a shape you can draw in one stroke. A contracting path downsamples the image the way a classifier would, to gather context. A symmetric expanding path upsamples back to full resolution, to recover location. And a set of skip connections hand the high-resolution detail from the contracting side directly across to the expanding side, so the network never has to reconstruct, from the squashed bottleneck, what it already saw sharply on the way down. Wrap that in a training recipe built for scarce data, and you can train the whole thing end to end from a few dozen images. We will build it up one piece at a time.

Two questions that fight each other

Before U-Net, the strongest way to segment a biomedical image was the sliding-window network of Ciresan and colleagues, the method that had won the 2012 version of this very challenge. The idea is direct: to label one pixel, cut out a square patch centered on it and ask a classifier what the center pixel is. Slide that window over every pixel and you get a full segmentation. It works, and it has two problems the paper is blunt about.

The first is speed. You run the network once per pixel, and neighboring patches overlap almost completely, so you recompute nearly the same features millions of times for a single image. The second is the tension from the last section, now made concrete as a single knob: the patch size. A big patch carries more context but needs more pooling, which blurs out exactly where the center pixel's boundary is. A small patch localizes sharply but sees too little around it to know what it is looking at. You are forced to trade the two against each other, and neither setting is good at both. Drag the patch size and feel the trade:

On top of that sits the data problem. The breakthrough classifiers of the era trained on a million labelled images. A lab annotating electron-microscopy stacks by hand might produce thirty. So the method also has to be frugal with labels. U-Net's answer to all three pressures, speed, the context-localization trade, and label scarcity, is one architecture and one training recipe. Start with the architecture.

The U: shrink to read it, grow it back

U-Net descends from the fully convolutional network (FCN), which had shown a few months earlier that you can turn a classifier into a dense predictor by replacing its pooling-driven downsampling with learned upsampling, and by combining a coarse, deep layer with a finer, shallower one to sharpen the result. U-Net takes that skeleton and makes it symmetric, deeper in the decoder, and trainable from almost no data. (One difference matters later: FCNadds the coarse and fine layers together. U-Net concatenates them, which is a bigger change than it sounds.)

Picture the network as the letter it is named for. Down the left arm runs the contracting path. It is an ordinary convnet: at each level, two $3\times3$ convolutions followed by ReLUs, then a $2\times2$ max-pool that halves the height and width. Every time the spatial size halves, the number of feature channels doubles, from 64 at the top to 1024 at the bottom. The image is getting smaller and "thicker": less space, more meaning per location. At the bottom of the U, the bottleneck, a tile that started at $572\times572$ has been squeezed to $28\times28$ with 1024 channels. This is the most context-rich, least spatially-precise the data ever gets. After four poolings each surviving cell has been shaped by a large patch of the original tile, and that accumulated reach is exactly what "context" means here; the price of the reach is that the same cell can no longer say where inside that patch anything actually sat.

That accumulated reach has a name: the receptive field, the patch of the original input that a single neuron has been shaped by. It grows by simple arithmetic as you descend. Each valid $3\times3$ convolution widens it by two pixels, and each $2\times2$ pool doubles the stride, so every later step reaches twice as far. Step the depth slider from the top to the bottleneck and watch the teal square, the receptive field of one neuron, swell from a 5-pixel sliver to a 140-pixel patch, about a quarter of the $572\times572$ tile, while the feature map on the right shrinks from $568$ to $28$. Deeper buys context and spends location:

Up the right arm runs the expanding path, and it undoes the squeeze. At each level an up-convolution doubles the height and width and halves the channels, two $3\times3$ convs refine the result, and the map climbs back toward full resolution. By the top the network is back to a large spatial map, now 64 channels deep, and a final $1\times1$ convolution turns each pixel's 64-number feature vector into a score for each class. Drag the marker down and back up the U and watch the shape of the data change at every step:

That is the whole forward pass, and it is just a feedforward convnet. Here it is in pseudocode. The only thing to notice is that the contracting path squirrels away each level's feature map in a list, and the expanding path pulls them back out in reverse. Those saved maps are the skip connections, and they are the reason the U works at all.

# U-Net forward pass: contract, then expand reusing the skips
skips = []
x = tile                              # one mirror-padded input tile
for ch in [64, 128, 256, 512]:        # contracting path
    x = relu(conv3x3(x)); x = relu(conv3x3(x))   # two valid convs
    skips.append(x)                   # keep this map for later
    x = maxpool2x2(x)                 # halve H,W; channels double next
x = relu(conv3x3(x)); x = relu(conv3x3(x))       # bottleneck, 1024 ch
for skip in reversed(skips):          # expanding path
    x = upconv2x2(x)                  # double H,W, halve channels
    x = concat(x, crop(skip, x))      # copy-and-crop the encoder map
    x = relu(conv3x3(x)); x = relu(conv3x3(x))
return conv1x1(x)                     # per-pixel class scores

A word on the up-convolution, because the name hides a choice. The paper describes it as "an upsampling of the feature map followed by a $2\times2$ convolution." Many later implementations collapse that into a single strided transposed convolution (also called a deconvolution, a misleading name since it is not the inverse of anything). The two are interchangeable in spirit, and people use the terms loosely, but they are not identical in practice: a strided transposed convolution can stamp a faint checkerboard into the output, which the upsample-then-convolve form avoids. The honest reading of the paper is the upsample-first one. What matters for the architecture is the bookkeeping it does: double the resolution, halve the channels.

The skip connections carry the detail across

The U-shape forces a question. By the time information reaches the bottleneck, it has been through four max-pools. A $28\times28$ map cannot, even in principle, say which of the original 572 columns a boundary fell on. The location has been quantized away. So if the decoder had only the bottleneck to work from, the best it could do is upsample a blurry, blocky guess: it would know roughly where the cell is, but it could not trace the cell's outline to the pixel. The fine detail is gone from that path. So the decoder needs the geometry handed back without losing the meaning it gathered: the decoder map should bring the meaning, a saved copy of the early features should bring the geometry, and the final convolutions should be free to weigh the two independently.

But the fine detail is not gone from the network. It was sitting right there in the contracting path's early layers, at full resolution, before the pooling destroyed it. The skip connection hands that lost detail back. At each level of the decoder, U-Net takes the saved feature map from the corresponding level of the contracting path, crops it to size, and concatenates it onto the upsampled decoder map before the next convolutions run. The decoder map brings the deep, wide-context meaning. The skip brings the sharp, high-resolution geometry. The convolutions that follow get to combine both, so the output can be both correctly classified and precisely placed.

This is the one place U-Net's choice of concatenation over FCN's addition earns its keep. Adding the two maps forces them into the same channels and lets the network blend them only in fixed proportion. Concatenating keeps every channel from both, so the following convolutions can learn for themselves how to weigh the precise-but-shallow skip against the meaningful-but-coarse decoder. It costs memory. It buys flexibility, and a symmetric decoder with as many channels as the encoder to use it.

You can feel what the skip does by taking it away. Below, the dashed amber curve is a cell's true outline. The teal mask is the network's prediction. With the skip turned off, the decoder works from the coarse bottleneck alone and produces a blocky, staircased boundary that misses every lobe. Turn the skip up and the high-resolution detail floods back in, and the predicted boundary snaps onto the true outline:

The skip connection is the whole idea. The rest, the unpadded convolutions, the tiling, the weighted loss, the deformations, is what it took to make this shape train well on real microscopy data with few labels.

Valid convolutions, and an image too big to fit

U-Net uses only the valid part of every convolution. A valid (or unpadded) $3\times3$ conv does not invent pixels at the image edge to keep the output the same size; it produces an output that is one pixel smaller on each side. Two convs per level, across the whole U, and the losses add up: an input tile of $572\times572$ comes out as a $388\times388$ segmentation. The output is smaller than the input by a fixed border, on purpose.

Why accept that? Because the alternative, padding with zeros, fabricates a hard black edge that no real tissue ever had, and the convolutions near the border learn from that fiction. Valid convolutions guarantee the opposite: every pixel in the output was computed only from real image context, never from invented padding. The segmentation map contains exactly the pixels for which the full context was actually available.

This creates a logistics problem. If the output is always smaller than the input, how do you segment an image larger than what fits on the GPU, or label the pixels right at the image's own edge, where there is no surrounding context to feed in? U-Net's answer is the overlap-tile strategy. You predict the image in tiles. To produce one output tile you feed in a larger input tile that includes a border of surrounding context. Where that border runs off the edge of the actual image, the missing context is invented by mirroring: the image is reflected across its own boundary, so the network sees a plausible continuation with the same texture and statistics instead of a black void. Drag the tile to the image's edge and watch the missing context fill in as a reflection:

One small constraint falls out of all the halving: you have to choose an input tile size so that every $2\times2$ max-pool lands on a map with even height and width, or the downsampling would not divide cleanly. It is a detail, but it is why the example tile is $572$ and not a round number. Halving an odd dimension has no integer answer, so a $2\times2$ pool on a 37-wide map has nowhere to put the leftover column and is undefined, which is why the input size is picked to stay even through every level of the descent.

A loss that learns to draw the borders

The output is a stack of class scores at every pixel. To turn scores into probabilities, U-Net runs a softmax across the classes at each pixel position $\mathbf{x}$:

Here $a_k(\mathbf{x})$ is the network's score for class $k$ at pixel $\mathbf{x}$, and $p_k(\mathbf{x})$ is the resulting probability, close to 1 for the winning class and near 0 for the rest. Training then pushes the probability of the correct class toward 1 at every pixel using a weighted cross-entropy, the paper's energy function:

Read it piece by piece. $\ell(\mathbf{x})$ is the true label at pixel $\mathbf{x}$, so $p_{\ell(\mathbf{x})}(\mathbf{x})$ is the probability the network assigned to the right answer there. When that probability is 1 the pixel contributes nothing; when it drops, $-\log$ of it grows, and the loss rises. Sum over every pixel $\mathbf{x}$ in the image domain $\Omega$, with a per-pixel weight $w(\mathbf{x})$ we are about to define, and you have the quantity to minimize.

(A note on the sign, because the paper's printed Eq (1) drops it. As published, the formula reads $E = \sum w\,\log p$ with no leading minus. Since a probability is at most 1, its log is at most 0, so that expression is never positive, and to drive the right-class probability toward 1 you would have to maximize it. The cross-entropy loss you actually minimize carries the minus sign, as in (1) above. It is a slip of convention, not of method.)

The weight map $w(\mathbf{x})$ is where U-Net solves a problem specific to cells: telling apart two cells of the same class that are touching. If the network only had to say "cell" or "background," it could merge two kissing cells into one blob and pay almost nothing for it, since the thin wall between them is just a sliver of pixels. So U-Net precomputes, for each training image, a weight map that makes those sliver pixels expensive to get wrong:

$w_c$ is a baseline weight that balances how often each class appears, so a rare class is not drowned out. The interesting term is the second one. $d_1(\mathbf{x})$ is the distance from pixel $\mathbf{x}$ to the border of the nearest cell, and $d_2(\mathbf{x})$ the distance to the second-nearest. Out in open background, far from everything, both distances are large, the exponential is near 0, and the weight is just the baseline. But in the thin gap between two touching cells, both borders are close, so $d_1 + d_2$ is small, the exponential approaches 1, and the weight jumps by up to $w_0$. The paper sets $w_0 = 10$ and $\sigma \approx 5$ pixels, so the separating membrane between two cells can carry roughly ten times the weight of an ordinary pixel. The exponential is the switch that makes this sharp: a near-zero exponent sits right at $e^0 = 1$, so the bonus fires at full strength in the thin gaps, while a large $d_1 + d_2$ drives the exponent steeply negative and the bonus dies off to almost nothing out in open tissue. Drag the two cells together and watch the ridge of weight ignite in the gap:

Put the two equations together and the training step is just: run the tile through the U, softmax the scores, take the per-pixel cross-entropy against the labels, multiply by the precomputed weight map, and sum.

# pixelwise softmax + weighted cross-entropy (Eq 1),
# using the precomputed border weight map (Eq 2)
p  = softmax(scores, over=classes)    # p_k(x) at every pixel x
ce = -log(p[label[x], x])             # per-pixel cross-entropy
w  = w_c + w0 * exp(-(d1 + d2)**2 / (2 * sigma**2))   # w0=10
loss = sum(w * ce)                    # summed over the output map

Thirty images, and the deformations that multiply them

An architecture and a loss are not enough when you have thirty training images. A network this size would just memorize them. U-Net's third leg is aggressive data augmentation, and the part that matters most is elastic deformation: warping each training image (and its label map) with a smooth, random distortion, so the network sees a fresh, plausible variant every time. Flips and rotations help too, but they only move a rigid image around. Real tissue is not rigid. It stretches and folds locally, and a network that has only seen rigid copies stays fragile to the warps it will actually meet.

The recipe is small. Lay a coarse $3\times3$ grid of control points over the image and give each one a random displacement drawn from a Gaussian with a standard deviation of 10 pixels. Interpolate those few displacements smoothly across every pixel (the paper uses bicubic interpolation) and you get a flowing, tissue-like warp, not a jagged scramble. Apply the same warp to the image and its segmentation so they stay aligned. Because real tissue deforms this way, the augmented images stay realistic, and the network learns to be invariant to exactly the kind of variation it will meet. Press for a new deformation and feel how one labelled cell becomes an endless supply of them:

Two smaller training choices round it out, and both are answers to the same constraint: tiny data on a small GPU. First, the batch is a single image. The authors would rather spend their memory on a large input tile, for more context per forward pass, than on a big batch, so they push the batch down to one and lean on a high momentum of $0.99$ to compensate. With momentum that high, each update is effectively an average over the last hundred-odd gradients, which steadies the very noisy single-image steps. (The "hundred-odd" is the intuition $1/(1-0.99)=100$, not a number the paper states.) Mechanically each step blends the fresh gradient with a fading memory of the ones before it, an exponential moving average over roughly the last hundred updates, so the quirks of any single image wash out and only the direction the images agree on actually moves the weights. Second, the weights are initialized from a Gaussian with standard deviation $\sqrt{2/N}$, where $N$ is the number of inputs feeding each neuron (for a $3\times3$ conv over 64 channels, $N = 9 \cdot 64 = 576$). That specific $\sqrt{2/N}$ is He initialization, designed so that ReLU layers neither blow up nor fade out as the signal passes through; the factor of 2 is the correction for ReLU zeroing half its inputs.

What it actually did

U-Net was tested on three biomedical segmentation tasks, and the same architecture won all of them. On the ISBI 2012 challenge for segmenting neuronal structures in electron-microscopy stacks, trained on just 30 images, it took the top of the ranking by warping error, the challenge's primary metric, and beat the prior sliding-window network it was designed to replace. On the ISBI 2015 cell-tracking challenge it won both light-microscopy categories outright, reaching an intersection-over-union (IoU, the overlap between predicted and true masks) of 92% on one dataset and 77.5% on the other, against second-best scores of 83% and 46%. Toggle between the two challenges:

That last caveat is worth stating plainly, since the paper is careful about it. The EM challenge reports three different error metrics, and U-Net was first by warping error, the one the ranking is sorted on. By Rand error, a couple of entries scored lower, and they leaned on heavy dataset-specific post-processing to do it. U-Net ranked first on the primary metric with no pre- or post-processing, from a handful of images. That is what made it a clean, general method rather than a tuned pipeline.

And it was fast and cheap to train. The whole thing took about ten hours on a single NVIDIA Titan with 6 GB of memory, and segmenting a $512\times512$ image at inference took under a second. A simple method that matched or beat specialized pipelines from a few dozen labelled images, and ran fast: that is why the architecture spread far past biology.

U-Net is four ideas in one shape. Contract to gather context. Expand to recover location. Skip the high-resolution detail across the gap so the recovery is exact rather than a guess. And train it on scarce data with a weighted loss that cares about the borders and an elastic augmentation that turns thirty images into many. The encoder-decoder-with-skips it introduced is now the default backbone for dense prediction everywhere, from medical imaging to the denoising network at the heart of modern diffusion models. The network could see fine all along. What it kept discarding was where things were, and U-Net kept a copy.