Source-linked AI summary

UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation

Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, Jianming Liang

arXiv:1912.05074v2eess.IVcs.CVcs.LG

TL;DR

Medical image segmentation architectures face unknown optimal depth and restrictive same-scale skip-connection fusion. UNet++ embeds partially shared U-Nets of varying depths, redesigns skip connections, and consistently improves semantic and instance segmentation across diverse biomedical imaging tasks.

  • Problem

    Encoder-decoder segmentation models have unknown optimal depth and restrictive skip connections that fuse only same-scale encoder and decoder features.

  • Method

    UNet++ embeds partially shared U-Nets of varying depths, densely connects decoder nodes, applies deep supervision, and supports pruning for faster inference.

  • Results

    UNet++ consistently outperforms U-Net and UNet+ across six segmentation tasks and multiple backbones, with deep supervision and voting increasing IoU by up to 0.8 points.

  • Takeaways & Limitations

    UNet++ provides a backbone-agnostic architecture for improved semantic and instance segmentation across multiple biomedical imaging applications.

Abstract

from arXiv · show

The state-of-the-art models for medical image segmentation are variants of U-Net and fully convolutional networks (FCN). Despite their success, these models have two limitations: (1) their optimal depth is apriori unknown, requiring extensive architecture search or inefficient ensemble of models of varying depths; and (2) their skip connections impose an unnecessarily restrictive fusion scheme, forcing aggregation only at the same-scale feature maps of the encoder and decoder sub-networks. To overcome these two limitations, we propose UNet++, a new neural architecture for semantic and instance segmentation, by (1) alleviating the unknown network depth with an efficient ensemble of U-Nets of varying depths, which partially share an encoder and co-learn simultaneously using deep supervision; (2) redesigning skip connections to aggregate features of varying semantic scales at the decoder sub-networks, leading to a highly flexible feature fusion scheme; and (3) devising a pruning scheme to accelerate the inference speed of UNet++. We have evaluated UNet++ using six different medical image segmentation datasets, covering multiple imaging modalities such as computed tomography (CT), magnetic resonance imaging (MRI), and electron microscopy (EM), and demonstrating that (1) UNet++ consistently outperforms the baseline models for the task of semantic segmentation across different datasets and backbone architectures; (2) UNet++ enhances segmentation quality of varying-size objects -- an improvement over the fixed-depth U-Net; (3) Mask RCNN++ (Mask R-CNN with UNet++ design) outperforms the original Mask R-CNN for the task of instance segmentation; and (4) pruned UNet++ models achieve significant speedup while showing only modest performance degradation. Our implementation and pre-trained models are available at https://github.com/MrGiovanni/UNetPlusPlus.

I. INTRODUCTION · II. PROPOSED NETWORK ARCHITECTURE: UNET++

UNet++ addresses encoder-decoder segmentation limitations by embedding U-Nets of varying depths and redesigning skip connections for flexible feature fusion. Evaluations across six datasets and multiple backbones report improved semantic and instance segmentation, including varying-size objects and multiple imaging modalities.

  • I. INTRODUCTION: Encoder-decoder segmentation networks use skip connections to combine coarse semantic decoder features with fine-grained encoder features and recover object details.These connections also support instance-level segmentation by helping distinguish each desired object instance.
  • I. INTRODUCTION: The optimal encoder-decoder depth varies with task difficulty and available labeled data, making separately training and ensembling models of different depths inefficient.This motivates an architecture that can represent multiple depths within one model.
  • II. PROPOSED NETWORK ARCHITECTURE: UNET++: UNet++ embeds U-Nets of varying depths whose decoders are densely connected at the same resolution through redesigned skip connections.The architecture is presented as a general-purpose image segmentation network that addresses the limitations of conventional encoder-decoder designs.
  • II. PROPOSED NETWORK ARCHITECTURE: UNET++: The redesigned skip connections enable flexible decoder feature fusion instead of restricting aggregation to same-scale encoder and decoder feature maps.UNet++ combines feature maps from the constituents U-Nets at the same resolution.
  • I. INTRODUCTION: Across six segmentation datasets and multiple backbones of different depths, UNet++ with redesigned skip connections and deep supervision achieves higher semantic and instance segmentation performance.The evaluated imaging modalities include computed tomography, magnetic resonance imaging, and electron microscopy.
  • I. INTRODUCTION: UNet++ provides a built-in ensemble of U-Nets of varying depths, improving segmentation of varying-size objects over the fixed-depth U-Net.The constituent U-Nets partially share an encoder while retaining their own decoders.
  • I. INTRODUCTION: Simultaneously training the embedded multi-depth U-Nets stimulates collaborative learning and outperforms individually training isolated U-Nets with the same architecture.The contribution is reported as a benefit of jointly learning the constituent networks within UNet++.

A. Motivation behind the new architecture · B. Technical details · 1) Network connectivity:

UNet++ addresses dataset-dependent depth selection and restrictive skip connections by combining U-Nets of varying depths with dense, cross-scale decoder connectivity. Deep supervision exposes segmentation outputs at multiple nodes, enabling inference-time pruning while preserving flexible feature fusion.

  • A. Motivation behind the new architecture: U-Net depth is dataset-dependent: shallower U-Net L3 performs better for Cell and Brain Tumor, whereas deeper U-Nets perform better for EM.The comparison used three relatively small datasets: Cell, EM, and Brain Tumor.
  • A. Motivation behind the new architecture: Deep supervision provides segmentation outputs at multiple nodes X0,j, allowing inference to select architectures from UNet++ (L1) through UNet++ (L4).Selecting X0,4 performs no pruning, whereas selecting X0,1 produces the maximally pruned architecture.
  • A. Motivation behind the new architecture: Because exhaustive neural architecture search is computationally constrained, U-Nete ensembles U-Nets of varying depths within one unified structure.Each ensemble member has a separate loss function and the U-Nets partially share the encoder while retaining their own decoders.
  • A. Motivation behind the new architecture: U-Nete’s disconnected decoders prevent deeper U-Nets from supervising shallower decoders, while its conventional skip connections impose restrictive feature fusion.These limitations motivate connecting decoder nodes across the ensemble.
  • A. Motivation behind the new architecture: UNet+ removes the original skip connections and connects every two adjacent ensemble nodes, allowing gradients from deeper decoders to reach shallower decoders.The new connectivity also relaxes the restrictive skip-connection design by changing how each decoding node receives features.
  • B. Technical details: UNet++ denotes node outputs as x_i,j, where i indexes encoder down-sampling layers and j indexes convolution layers within dense skip-connection blocks.The feature-map stack is computed through convolution, activation, down-sampling, up-sampling, and concatenation operations.
  • 1) Network connectivity:: At j = 0, nodes receive one encoder input; at j = 1, they receive two encoder inputs; and at j > 1, they receive j + 1 inputs including a lower skip-connection output.Dense convolution blocks cause prior feature maps to accumulate at the current node.

2) Deep supervision:

UNet++ applies deep supervision at multiple semantic scales by attaching class-specific prediction heads and optimizing each decoder with a hybrid cross-entropy and soft Dice loss. The overall objective equally weights the losses from all decoders without Gaussian-blurring the supervision targets.

  • Deep supervision: Deep supervision attaches a 1×1 convolution with C kernels and a Sigmoid activation to nodes X0,1, X0,2, X0,3, and X0,4.C denotes the number of classes in the dataset.
  • Deep supervision: Each semantic scale uses a hybrid segmentation loss combining pixel-wise cross-entropy and soft Dice-coefficient losses.The formulation uses target labels and predicted class probabilities for pixels in each batch.
  • Deep supervision: The overall UNet++ objective is a weighted sum of the hybrid losses from the individual decoders.The decoder index d determines which supervised output contributes to the total loss.
  • Deep supervision: Experiments assign equal balanced weights to every decoder loss, with ηi ≡1, and do not apply Gaussian blur to ground-truth supervision targets.The same weighting is used across decoder outputs.

3) Model pruning: … 2) Cell:

The study combines deep-supervision-enabled UNet++ pruning with evaluation across six biomedical segmentation datasets and detailed EM and Cell subsets. Pruning selects one segmentation branch for speed gains, while datasets span microscopy, CT, and MRI modalities.

  • 3) Model pruning:: Deep supervision supports ensemble and pruned UNet++ operation modes, with pruning determined by selecting a single segmentation branch.Ensemble mode averages outputs from all segmentation branches; pruned mode uses one branch, controlling the extent of pruning and speed gain.
  • A. Datasets: The study evaluates six biomedical image segmentation datasets spanning microscopy, computed tomography (CT), and magnetic resonance imaging (MRI).The datasets cover lesions and organs from commonly used medical imaging modalities.
  • A. Datasets: Figure 3 compares U-Net, wide U-Net, and UNet++ across six biomedical image segmentation applications using 2D and 3D modalities.Quantitative prediction scores are reported as IoU | Dice.
  • 1) Electron Microscopy (EM):: The EM segmentation challenge dataset contains 30 images of 512×512 pixels from serial-section transmission electron microscopy of the Drosophila first-instar larva ventral nerve cord.Each image includes fully annotated segmentation maps for cells and membranes.
  • 1) Electron Microscopy (EM):: The EM dataset’s labeled images are divided into training, validation, and test subsets.The supplied passage states the split but does not provide all subset counts.
  • 2) Cell:: The Cell dataset was acquired with a Cell-CT imaging system and contains two binary cell masks per image from two trained expert annotators.Experiments select 354 images with the highest annotator agreement.
  • 2) Cell:: The selected Cell images are split into training (212 images), validation (70 images), and test (72 images) subsets.The split applies to the 354-image subset selected for highest agreement between the two expert annotators.

3) Nuclei: … B. Baselines and implementation

The study evaluates segmentation across nuclei, brain tumor, liver, and lung-nodule datasets, using standard and parameter-matched wide U-Net/V-Net baselines alongside UNet+ and UNet++. Experiments use dataset-specific preprocessing and splits, early stopping, Dice and IoU evaluation, and Adam optimization.

  • 3) Nuclei:: The nuclei dataset contains 670 segmented images from brightfield and fluorescence modalities, with instance-level nucleolus annotations.Images are randomly divided into training (50%), validation (20%), and test (30%) sets; 96×96 patches are extracted using a sliding window.
  • 4) Brain Tumor:: The brain-tumor experiments use 66,348 MR slices from 30 BraTS 2013 patients across High-grade and Low-grade cases.The scans include Flair, T1, T1c, and T2 modalities, are rescaled to 256×256, and are randomly assigned into five patient-level folds.
  • 5) Liver:: The liver dataset comprises 331 CT scans, split into training (100 patients), validation (15 patients), and test (15 patients) subsets.Experiments treat liver as the positive class and all other regions as negative.
  • 6) Lung Nodule:: The lung-nodule dataset contains 1,018 cases from LIDC-IDRI, with six cases removed for ground-truth issues.The remaining cases are split into training (510), validation (100), and test (408) sets; nodules are represented as volumetric binary masks in 3D CT scans.
  • B. Baselines and implementation: The study compares UNet++ with original U-Net and parameter-matched wide U-Net for 2D tasks, and V-Net and wide V-Net for 3D tasks.U-Net and V-Net provide common performance baselines, while wider variants are designed with similar parameter counts to the proposed architectures.
  • B. Baselines and implementation: UNet++ is additionally compared with UNet+, an intermediate architecture whose intermediate-node kernel counts are specified in Table III.Table III details the architectures and notes that wider U-Net and V-Net variants have comparable parameter counts to UNet++ and VNet++.
  • B. Baselines and implementation: Experiments use Keras with a TensorFlow backend, early stopping on the validation set, Dice coefficient and IoU evaluation, and Adam with a 3e-4 learning rate.Pixel-wise sensitivity, specificity, F1, and F2 scores, together with statistical analysis, are reported in Appendix Section A.

IV. RESULTS · A. Semantic segmentation results

UNet++ consistently improves semantic segmentation over U-Net and UNet+ across datasets, applications, and backbone architectures. Its gains are also demonstrated against wide U-Net and in the BraTS 2013 brain-tumor comparison.

  • A. Semantic segmentation results: UNet++ without deep supervision achieves significant IoU gains over U-Net and wide U-Net across all six segmentation tasks.The reported gains include 0.62±0.10 versus U-Net and 0.55±0.01 versus wide U-Net for neuronal structure, and 2.30±0.30 for cell segmentation versus U-Net.
  • A. Semantic segmentation results: Wide U-Net consistently outperforms U-Net, which the paper attributes to wide U-Net’s larger number of parameters.The comparison reports both parameter counts and segmentation results for the six tasks.
  • A. Semantic segmentation results: The semantic-segmentation comparison covers neuronal structure, cell, nuclei, brain tumor, and liver applications across modern CNN backbones.The backbone study applies redesigned UNet++ skip connections to vgg-19, resnet-152, and densenet-201, while also training U-Net and UNet+ for comparison.
  • A. Semantic segmentation results: UNet++ is evaluated against U-Net, wide U-Net, and UNet+ using IoU (mean±s.d. %) across six segmentation tasks.UNet+ and UNet++ are tested with and without deep supervision, with statistical comparisons performed over 20 independent trials.
  • A. Semantic segmentation results: UNet++ with deep supervision reaches 90.83%±2.46% complete Dice on BraTS 2013, compared with the winner’s 92%.The paper notes that the BraTS 2013 winner holds a “complete” Dice of 92%.
  • A. Semantic segmentation results: UNet++ with deep supervision consistently outperforms U-Net across all evaluated backbone architectures and applications.Dense connections between intermediate layers also produce higher segmentation performance than UNet+ in most experimental configurations.
  • A. Semantic segmentation results: UNet++ consistently outperforms U-Net and UNet+ across EM, Cell, Nuclei, Brain Tumor, and Liver segmentation datasets.The comparison uses vgg-19, resnet-152, and densenet-201 backbones converted into U-Net models with decoder sub-networks and redesigned skip connections.

B. Instance segmentation results · C. Model pruning

UNet++ redesigns Mask R-CNN skip connections for instance segmentation, producing Mask RCNN++ for nuclei segmentation. Its deep supervision also enables pruning decoder paths to reduce inference time, with performance evaluated to select the pruning level.

  • B. Instance segmentation results: Instance segmentation is more challenging than semantic segmentation because it must segment and distinguish all object instances.Mask R-CNN is used as the baseline model for this task.
  • B. Instance segmentation results: Mask RCNN++ replaces FPN’s plain skip connections with UNet++ redesigned skip connections while retaining Mask R-CNN for instance segmentation.The model uses ResNet-101 as its backbone.
  • B. Instance segmentation results: The Nuclei dataset supports both semantic segmentation of all nuclei as foreground and instance segmentation of each nucleus separately.Distinct nucleolus instances are annotated in different colors, making them separate objects.
  • B. Instance segmentation results: Mask RCNN++ outperforms Mask R-CNN for nuclei instance segmentation.The comparison is reported in Table V, which evaluates the redesigned skip connections for instance segmentation.
  • C. Model pruning: Deep supervision makes each UNet++ decoder path independent at inference, allowing the decoder at depth d+1 to be removed to obtain a shallower depth-d model.This pruning procedure is possible after training because the decoder paths are completely independent.
  • C. Model pruning: Pruning can significantly reduce inference time, but segmentation performance may degrade.The pruning level should therefore be selected by evaluating model performance.
  • C. Model pruning: Embedded shallower networks show improved segmentation after pruning from UNet++ compared with the same networks trained in isolation.UNet++ L4 has identical isolated and embedded performance because it undergoes no pruning.

D. Embedded vs. isolated training of pruned models … B. Feature maps visualization

UNet++ supports embedded training of pruned subnetworks and outperforms isolated training, especially under aggressive pruning. Its dense, multiscale design improves segmentation across tumor sizes and promotes feature fusion and optimization.

  • D. Embedded vs. isolated training of pruned models: Embedded training jointly trains the full UNet++ graph with deep supervision, then uses the selected yellow subnetwork for inference after pruning.Isolated training instead trains UNet++ L_d independently, without interactions with deeper encoder and decoder nodes.
  • V. DISCUSSIONS: The UNet++ design is discussed as an efficient ensemble of U-Nets with varying depths and multiscale receptive fields.This ensemble addresses segmentation across different object sizes while enabling pruned subnetworks to share the encoder.
  • D. Embedded vs. isolated training of pruned models: Embedded training produces higher-performing pruned models than isolated training across pruning levels and datasets.The advantage is more pronounced under aggressive pruning, particularly when the full model is reduced to UNet++ L1.
  • A. Performance analysis on stratified lesion sizes: UNet++ consistently outperforms U-Net for brain tumors grouped into seven size buckets.The comparison evaluates segmentation performance separately across tumor-size groups.
  • A. Performance analysis on stratified lesion sizes: Five of seven tumor-size comparisons are statistically significant at p < 0.05.The improvement is attributed to UNet++’s built-in ensemble of U-Nets, which supports segmentation using multi-receptive-field networks.
  • B. Feature maps visualization: Redesigned skip connections fuse semantically rich decoder maps with intermediate maps spanning varying semantic scales.Feature-map visualization is used to illustrate this multiscale fusion capability.
  • B. Feature maps visualization: Feature-map visualizations compare representative early, intermediate, and late activations along the topmost skip connection for a brain-tumor image.Representative maps are formed by averaging all feature maps at each layer, under architectures with different supervision configurations.
  • B. Feature maps visualization: Dense connections in UNet++ encourage better optimization and reach lower validation loss across six medical applications.The learning curves are presented in the appendix.

C. Collaborative learning in UNet++ … B. Feature aggregation

UNet++ combines collaborative learning with redesigned skip connections and multiscale feature aggregation. The related-work discussion situates these components among FCN, U-Net, GridNet, and FRRN-based architectures.

  • C. Collaborative learning in UNet++: UNet++ embodies collaborative learning by aggregating multiple-depth networks and supervising segmentation heads from each constituent network.Collaborative learning trains multiple classifier heads simultaneously on the same data and is reported to improve generalization.
  • C. Collaborative learning in UNet++: UNet++ segmentation heads receive gradients from both strong ground-truth losses and additional supervision signals.The supplied passage gives X0,2 as an example of a segmentation head receiving these gradients, but its description is truncated.
  • VI. RELATED WORKS: The related-work discussion identifies redesigned skip connections, feature aggregation, and deep supervision as the main components of UNet++.These components organize the review of prior work.
  • A. Skip connections: FCN introduced skip connections for semantic segmentation, while U-Net subsequently applied skip connections to medical-image segmentation.The passage attributes FCN to Long et al. and U-Net to Ronneberger et al.
  • A. Skip connections: FCN and U-Net differ in how up-sampled decoder feature maps are fused with same-scale encoder feature maps.The supplied passage begins describing FCN’s summation-based fusion, but the comparison is truncated.
  • B. Feature aggregation: GridNet explores hierarchical feature aggregation through an encoder-decoder architecture whose feature maps are wired in a grid fashion.The passage presents GridNet as generalizing several classical segmentation architectures.
  • B. Feature aggregation: GridNet does not represent UNet++ because it lacks up-sampling layers between skip connections.The passage also notes that GridNet contains multiple streams with different resolutions.
  • B. Feature aggregation: FRRN uses two streams, carrying full-resolution information in one stream and context information in a separate pooling stream.The related work also cites incremental MRRN with 28.6M parameters and dense MRRN with 25.5M parameters as improved FRRN versions.

C. Deep supervision · D. Our previous work · VII. CONCLUSION

The paper situates UNet++ within deep-supervision research, extends its initial formulation through broader studies and applications, and attributes improved segmentation performance to its nested structure and redesigned skip connections.

  • C. Deep supervision: Deep supervision improves hidden-layer learning by encouraging intermediate layers to learn discriminative features, supporting faster convergence and network regularization.Prior work also uses deep supervision in U-Net-like architectures and implicitly in DenseNet.
  • C. Deep supervision: Combining predictions from feature maps at varying resolutions can address optimization difficulties while improving convergence and discriminative capability.Related architectures added multiple deeply supervised layers, whereas UNet++ uses multiple decoders that automatically generate full-resolution segmentation maps.
  • D. Our previous work: UNet++ was first presented in DLMIA 2018 and was subsequently adopted as a strong baseline, architectural inspiration, and tool across biomedical, natural, and satellite-image segmentation.The supplied passage lists adoption by the research community and applications across these image domains.
  • D. Our previous work: The current work extends the prior UNet++ study with a comprehensive network-depth analysis and evidence that embedded multi-depth U-Nets outperform individually trained isolated models.The comparison concerns embedded versus isolated training at various levels of pruned UNet++.
  • D. Our previous work: Further extensions add an MRI brain-tumor dataset, demonstrate Mask RCNN++, examine advanced encoder backbones, and study segmentation of lesions with varying sizes.These extensions strengthen the experiments and test UNet++ across additional tasks and model configurations.
  • VII. CONCLUSION: UNet++ is presented as a more accurate segmentation architecture whose improved performance is attributed to its nested structure and redesigned skip connections.The design addresses the U-Net’s unknown optimal depth and unnecessarily restrictive skip-connection design.

APPENDIX A ADDITIONAL MEASUREMENTS · APPENDIX B LEARNING CURVES

Appendix measurements show that UNet++ with deep supervision significantly outperforms U-Net across six biomedical applications. Learning curves further show faster convergence and lower validation loss across six segmentation tasks.

  • APPENDIX A ADDITIONAL MEASUREMENTS: UNet++ with deep supervision achieves significantly higher segmentation performance than the original U-Net across all six biomedical applications.The comparison uses pixel-wise sensitivity, specificity, F1, and F2 scores, with p-values calculated between the two models.
  • APPENDIX A ADDITIONAL MEASUREMENTS: The additional measurements report pixel-wise sensitivity for all six applications under study.Sensitivity is presented alongside specificity, F1, and F2 scores.
  • APPENDIX A ADDITIONAL MEASUREMENTS: The additional measurements report pixel-wise specificity, F1, and F2 scores for all six applications under study.These metrics complement the reported pixel-wise sensitivity results.
  • APPENDIX B LEARNING CURVES: Learning curves compare UNet++ and U-Net on neuronal structure, cell, nuclei, brain tumor, liver, and lung nodule segmentation.Validation losses were averaged over 20 trials for each application.
  • APPENDIX B LEARNING CURVES: UNet++ with deep supervision accelerates convergence relative to U-Net across the evaluated segmentation tasks.The learning curves attribute this optimization advantage to the redesigned intermediate layers and dense skip connections.
  • APPENDIX B LEARNING CURVES: UNet++ with deep supervision yields lower validation loss than U-Net across the evaluated segmentation tasks.The lower loss is evidenced by validation learning curves for six applications.
Loading 1912.05074v2…