U-Net: Convolutional Networks for Biomedical Image Segmentation

“U-Net: Convolutional Networks for Biomedical Image Segmentation” by Olaf Ronneberger, Philipp Fischer and Thomas Brox of the University of Freiburg was posted to arXiv on 18 May 2015 and published at Medical Image Computing and Computer-Assisted Intervention (MICCAI) 2015, in Lecture Notes in Computer Science. In eight pages it presented a fully convolutional network with a contracting path for context and a symmetric expanding path for localization, and a training strategy built on heavy data augmentation, and showed that the pair could learn pixel-accurate segmentation of microscopy images from a few dozen annotated examples. The network as understood today is covered in the U-Net article. The Springer chapter text could not be consulted for this article, so details below follow arXiv v1.

Problem

The paper opens with the success of deep convolutional networks on classification, where an image gets one label, and points to two ways biomedical tasks differ: the output must be localized, with a class for every pixel, and thousands of annotated training images are usually out of reach.

The method to beat was the network of Cireşan et al. (2012), which won the ISBI 2012 electron-microscopy (EM) segmentation challenge by classifying each pixel from a square patch centered on it. Ronneberger et al. credit it with two strengths: it localizes, and a few images yield many training patches. They name two drawbacks. It is slow, because the network runs separately for every patch, and overlapping patches make much of the computation redundant. And it trades localization against context: larger patches need more max pooling, which costs localization accuracy, while small patches show the network little context. The authors note that classifiers combining features from several layers, such as hypercolumns (Hariharan et al., 2015), had begun to show that good localization and context are possible together.

They took as their starting point the fully convolutional network of Long, Shelhamer and Darrell, then an arXiv preprint, which predicts every pixel in one pass and combines upsampled coarse outputs with high-resolution features from the contracting network. Their goal was to make that design work with very few training images and give more precise segmentations.

Contribution

The introduction sets out four ideas:

• A symmetric expanding path. FCN upsampled class scores; U-Net’s upsampling part keeps a large number of feature channels, so context can propagate to the high-resolution layers. The expanding path therefore mirrors the contracting path, and the network takes the U shape that gives it its name. Features from the contracting path are cropped and concatenated at every scale (encoder–decoder covers the general design).
• Strong data augmentation with elastic deformations. The authors argue that deformation is the most common variation in tissue and can be simulated realistically, so random elastic deformations teach invariance the small training sets cannot show. They suggest it is the key to training with very few annotated images.
• A weighted loss for touching objects. Background pixels in the thin gaps between touching cells receive a large weight in a pixel-wise cross-entropy, so that a semantic segmentation also keeps instances apart.
• The overlap-tile strategy. Because the network uses only the valid part of each convolution, every output pixel has its full context in the input; large images are processed as tiles whose inputs overlap, with context beyond the image border filled by mirroring. Without tiling, the authors note, GPU memory would limit the image resolution.

Method

The method section is short: two paragraphs on the architecture, which has 23 convolutional layers, and two pages on training. The training text gives the choices that made the method practical on a 2015 GPU. Large tiles were preferred to large batches, so a batch held a single image, compensated by a momentum of 0.99. Weights were drawn from a Gaussian with standard deviation $\sqrt{2/N}$, following the initialization He et al. (2015) derived for rectifier networks, which the authors call extremely important for a deep network with several paths. Dropout at the end of the contracting path is described as further implicit augmentation. The equations, layer sizes and augmentation parameters are worked through in the U-Net article.

Results

The experiments use three datasets, all microscopy, and the same network throughout.

ISBI 2012 EM segmentation. The challenge, still open to submissions in 2015, provides 30 annotated $512 \times 512$ serial-section transmission EM images of the Drosophila first instar larva ventral nerve cord; test labels are secret and the organizers score submitted membrane probability maps by warping, Rand and pixel error. Averaging predictions over seven rotated inputs, with no further pre- or post-processing, U-Net reached a warping error of 0.0003529, a new best, and a Rand error of 0.0382, against 0.000420 and 0.0504 for the best submission of Cireşan et al. Table 1, the leaderboard of 6 March 2015 sorted by warping error, puts U-Net first, narrowly ahead of an entry named DIVE-SCI (0.000355). The table also shows entries with lower Rand error; the authors attribute these to highly dataset-specific post-processing applied to Cireşan et al.’s probability map, and note in a footnote that one such group submitted 78 solutions.

ISBI cell tracking challenge 2015. On the segmentation part of the challenge (Maška et al., 2014), the authors report the mean IoU on two transmitted-light datasets with partially annotated training sets: PhC-U373, phase-contrast images of glioblastoma-astrocytoma cells with 35 training images, and DIC-HeLa, with 20. U-Net scored 0.9203 on PhC-U373 against 0.83 for the second-best 2015 entry, and 0.7756 on DIC-HeLa against 0.46, margins of about 9 and 32 points. On these two categories the authors won the challenge.

Cost. Training took 10 hours on an NVIDIA Titan GPU with 6 GB, and segmenting a $512 \times 512$ image took less than a second on a recent GPU.

Code. The paper announced the full Caffe-based implementation and the trained networks. The project page at the University of Freiburg still hosts a release archive dated 2 October 2015, containing the trained network, a modified Caffe, a MATLAB interface for overlap-tile segmentation and the greedy tracking algorithm used for the cell tracking submission. The page also reports a win in the ISBI 2015 challenge on caries detection in bitewing radiographs, which the paper does not include.

Small inconsistencies. Eq. 1, the weighted cross-entropy, is printed without its minus sign in arXiv v1. The DIC-HeLa IoU appears as 77.5% in the text and 0.7756 in Table 2, which would round to 77.6%. Neither changes a conclusion.

Impact

U-Net became a standard architecture for biomedical segmentation, and the paper is widely cited. Its own authors extended it first: 3D U-Net (Çiçek et al., 2016), with Ronneberger and Brox as co-authors, replaced the 2D operations by 3D ones and trained from sparsely annotated volumes, and its first author received the MICCAI Young Scientist Publication Impact Award in 2021 for the 3D U-Net paper. nnU-Net (Isensee et al., 2021) kept the U-Net and automated the decisions around it, from preprocessing and topology to training and post-processing; without manual intervention it surpassed most existing approaches on 23 public datasets from biomedical segmentation competitions.

The design also left medicine. pix2pix (Isola et al., 2017) used a U-Net generator for image-to-image translation, the skips letting low-level structure pass straight from input to output, and denoising diffusion probabilistic models (Ho et al., 2020) used a U-Net backbone with self-attention for their reverse process.

Limitations

The paper has no limitations section; the points below are this article’s reading of it.

• No ablations. The paper does not isolate the contribution of the symmetric decoder, the elastic deformations or the weight map; the claim that elastic deformation is the key concept is argued, not measured.
• A narrow evaluation. All three datasets are 2D microscopy with 20 to 35 training images each, and Table 2 compares U-Net on the cell tracking data with three 2014 entries and only the second-best 2015 score. On the EM leaderboard U-Net led on warping error by a small margin and, by Table 1, trailed other entries on Rand and pixel error.
• Underspecified details. The paper does not say how the upsampling before each “up-convolution” is computed, how the class-balancing weights $w_c$ are set, or how instances were extracted for the cell tracking evaluation.
• Constraints from valid convolutions. Input tiles must have sizes for which every pooling step sees an even size, outputs are 184 pixels narrower than inputs, and large images need tiling with mirrored borders; many later implementations pad instead.
• Instance separation through labels. Keeping touching cells apart depends on carving gaps into the training masks and weighting them, which requires instance annotations and fails where no gap exists.

What Came After

The volumetric extensions (3D U-Net, V-Net), attention-gated and nested-skip variants (Attention U-Net, UNet++) and nnU-Net’s self-configuring pipeline carried the architecture through medical imaging; the U-Net article covers them. In natural-image semantic segmentation, FCN’s successors such as DeepLab relied on dilated convolutions and little or light decoding, while encoder–decoders with concatenated skips became a common choice wherever an output must align pixel for pixel with the input, in image segmentation and in image generation.