PaperlayerView exact revision on arXiv

Source-linked AI summary

UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation

Abdelrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, Fahad Shahbaz Khan

arXiv:2212.04497v3cs.CV

TL;DR

Volumetric medical segmentation is challenged by the quadratic cost of self-attention. UNETR++ uses efficient paired spatial and channel attention, achieving favorable results across five datasets while reducing model complexity and improving speed.

  • Problem

    Quadratic self-attention makes transformer-based volumetric medical segmentation computationally expensive, while existing hybrid methods increase model size and FLOPs.

  • Method

    UNETR++ is a hierarchical 3D segmentation approach using an efficient paired-attention block with interdependent spatial and channel branches and shared query-key mappings.

  • Results

    UNETR++ achieves favorable segmentation results on five datasets; on Synapse, it attains a DSC of 87.22%.

  • Takeaways & Limitations

    UNETR++ combines competitive segmentation accuracy with lower model complexity, memory consumption, and inference speed than existing methods.

  • Takeaways & Limitations

    UNETR++ and existing methods struggle to segment certain organs in a few Synapse outlier cases.

Abstract

from arXiv · show

Owing to the success of transformer models, recent works study their applicability in 3D medical segmentation tasks. Within the transformer models, the self-attention mechanism is one of the main building blocks that strives to capture long-range dependencies. However, the self-attention operation has quadratic complexity which proves to be a computational bottleneck, especially in volumetric medical imaging, where the inputs are 3D with numerous slices. In this paper, we propose a 3D medical image segmentation approach, named UNETR++, that offers both high-quality segmentation masks as well as efficiency in terms of parameters, compute cost, and inference speed. The core of our design is the introduction of a novel efficient paired attention (EPA) block that efficiently learns spatial and channel-wise discriminative features using a pair of inter-dependent branches based on spatial and channel attention. Our spatial attention formulation is efficient having linear complexity with respect to the input sequence length. To enable communication between spatial and channel-focused branches, we share the weights of query and key mapping functions that provide a complimentary benefit (paired attention), while also reducing the overall network parameters. Our extensive evaluations on five benchmarks, Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung, reveal the effectiveness of our contributions in terms of both efficiency and accuracy. On Synapse, our UNETR++ sets a new state-of-the-art with a Dice Score of 87.2%, while being significantly efficient with a reduction of over 71% in terms of both parameters and FLOPs, compared to the best method in the literature. Code: https://github.com/Amshaker/unetr_plus_plus.

1. Introduction

UNETR++ addresses the efficiency–accuracy challenge in volumetric medical image segmentation by introducing an efficient paired-attention design within a U-Net-like framework. The approach combines spatial and channel attention and is evaluated on five medical segmentation benchmarks.

  • Background: Volumetric 3D segmentation supports applications such as tumor identification and organ localization for diagnosis.
  • Motivation: Hybrid convolution-and-attention architectures improve segmentation accuracy but substantially increase model parameters and FLOPs.
  • Method: UNETR++ introduces an efficient paired-attention block with spatial and channel attention branches.The spatial branch projects keys and values into a fixed lower-dimensional space, making self-attention linear in the number of input tokens; the channel branch models dependencies between channel feature maps.
  • Evaluation: UNETR++ is evaluated on five benchmarks: Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lungs.The experiments include both qualitative and quantitative comparisons focused on segmentation accuracy and model efficiency.

2. Related Work

Related work spans CNN-based, transformer-based, and hybrid architectures for medical image segmentation. Hybrid methods improve segmentation performance over pure CNNs and transformer designs, but may substantially increase parameters and FLOPs.

  • CNN-based Segmentation Methods: CNN-based approaches extend U-Net architectures for diverse medical image segmentation tasks, including 3D volumetric segmentation.The supplied passage identifies U-Net as the foundational design and notes multiple extensions for 3D medical image segmentation.
  • Transformers-based Segmentation Methods: Transformer-based methods use self-attention to model interactions among image-patch sequences and learn global relationships.Vision transformers have gained popularity in vision tasks such as classification and detection because they encode long-range dependencies.
  • Transformers-based Segmentation Methods: Recent transformer works investigate reducing the complexity of standard self-attention within transformer frameworks.The passage lists several efforts focused on alleviating self-attention’s complexity issue.
  • Hybrid Segmentation Methods: Hybrid architectures combine convolution and self-attention for segmentation, including TransFuse’s parallel CNN-transformer design and MedT’s gated position-sensitive axial attention.TransFuse fuses multi-level encoder features through BiFusion, while MedT controls positional embedding in self-attention with gated axial attention.
  • Our Approach: Hybrid methods such as UNETR and nnFormer improve segmentation performance over pure CNNs and transformer-based counterparts but create substantially larger models.UNETR is cited as having 2.5× more parameters, though the supplied passage truncates the comparison before stating the baseline’s identity.

3. Method

UNETR++ uses a hierarchical encoder-decoder with skip connections and convolutional blocks, centered on an efficient paired-attention (EPA) block. EPA jointly models spatial and channel features through shared query-key mappings, while spatial attention reduces self-attention complexity from quadratic to linear; training combines soft Dice and cross-entropy losses.

  • UNETR++ architecture: Skip connections link encoder and decoder stages to merge multiresolution features and recover spatial information lost during downsampling.The final decoder features are fused with convolutional feature maps to enhance representation before producing the segmentation mask.
  • UNETR++ architecture: UNETR++ employs a four-stage hierarchical encoder and decoder, progressively changing feature resolution through downsampling and upsampling by factors of two.The architecture uses 3D patch embedding in the first encoder stage, deconvolution-based upsampling in decoder stages, and reduces decoder channels by two between stages.
  • Efficient paired attention: EPA combines spatial and channel attention modules with shared query-key weights and distinct value layers to learn complementary spatial-channel representations.The two branches receive the same input feature maps, and their outputs are fused and transformed with convolution blocks.
  • Efficient paired attention: O(np) replaces O(n^2) in the spatial attention module, where n is the token count and p is the vector dimension.Channel attention separately captures interdependencies between feature channels using dot-product operations in the channel dimension.
  • Loss function: The training objective sums soft Dice loss and cross-entropy loss to leverage their complementary benefits across classes and voxels.The formulation uses ground-truth labels and predicted class probabilities at each voxel.

4. Experiments

Experiments evaluate UNETR++ across five 3D medical segmentation benchmarks using controlled training and testing protocols. The results show reduced complexity with competitive or improved segmentation performance, including strong results on Synapse, BTCV, BraTS, and lung segmentation.

  • Datasets and Protocol: Experiments cover Synapse, BTCV, ACDC, BraTS, and Medical Segmentation Decathlon-Lung datasets with task-specific segmentation targets.Synapse contains 30 abdominal CT subjects with eight organs; 18 samples are used for training and 12 for evaluation.
  • Synapse Dataset: 87.22% DSC is achieved on Synapse when EPA blocks are introduced in both encoders and decoders, improving over the encoder-only result of 85.17%.The comparison reports single-model accuracy without pre-training, ensembles, or additional data.
  • Ablation and Complexity: UNETR++ reduces parameters from 92.49M to 16.60M and FLOPs from 75.76G to 30.75G while maintaining a DSC of 78.29% against baseline UNETR.The hierarchical architecture uses four encoder and four decoder stages with feature-map downsampling after each encoder stage.
  • Additional Benchmarks: 92.83% mean DSC is achieved on the reported benchmark, exceeding UNETR at 86.61% and nnFormer at 92.06%.The cited passage introduces this result immediately before the BraTS evaluation and does not identify the benchmark in the excerpt.
  • Additional Benchmarks: 80.68% mean DSC is achieved on the lung cancer segmentation task, outperforming existing methods.BRaTS experiments additionally compare segmentation performance, model complexity, and inference time under a fixed input size and preprocessing strategy.

5. Conclusion

UNETR++ is a hierarchical approach for 3D medical segmentation that introduces an efficient paired attention block to encode inter-dependent spatial and channel features. Sharing query and key mapping weights improves communication between branches while reducing parameters.

  • 5. Conclusion: UNETR++ introduces an efficient paired attention block for 3D medical segmentation.The block encodes enriched inter-dependent spatial and channel features using spatial and channel attention.
  • 5. Conclusion: Sharing query and key mapping weights improves communication between spatial and channel branches while reducing parameters.The shared mappings provide complementary benefits within the paired attention design.

A. Additional Implementation Details

UNETR++ uses a four-stage hierarchical encoder-decoder with three EPA blocks per stage, while the final decoder stage replaces EPA processing with a 3×3×3 convolution to reduce computation at large spatial sizes.

  • Overall Architecture: The encoder has four stages with channel counts [32, 64, 128, 256], each containing three four-head EPA blocks.The channel counts correspond to stages [C1, C2, C3, C4].
  • Overall Architecture: The decoder has four stages, each using factor-of-two deconvolutional upsampling followed by three EPA blocks.Deconvolution increases feature-map resolution by a factor of two.
  • Overall Architecture: At the final decoder stage, a 3×3×3 convolutional block replaces heavy self-attention because the spatial size becomes substantially larger, reaching [128, 128, 64, 16] for Synapse.This design compensates for the computational cost of self-attention at the largest spatial resolution.

B. Additional Qualitative Results

The section presents additional qualitative comparisons involving UNETR++ on the Synapse and ACDC datasets. It also details comparisons with a baseline across Synapse, ACDC, and Dechatlon-Lungs.

  • Qualitative comparisons: UNETR++ is qualitatively compared with state-of-the-art methods on Synapse and ACDC.These comparisons are presented as additional qualitative results.
  • Baseline comparison: Detailed comparisons between UNETR++ and a baseline cover the Synapse, ACDC, and Dechatlon-Lungs datasets.The baseline analysis spans all three named datasets.
  • Comparison scope: The section combines comparisons against state-of-the-art methods with a separate baseline comparison.Both comparison types are reported for the additional qualitative analysis.

B.1. Synapse Dataset

On Synapse, qualitative comparisons show that UNETR++ produces more accurate segmentation predictions across organ sizes and scales than existing approaches. It distinguishes challenging tissues, segments all organs more completely, and avoids several under- and over-segmentation errors observed in competing methods.

  • UNETR++ differentiates stomach tissues at different sizes and successfully segments the spleen in the first two qualitative cases.
  • UNETR++ provides better multi-scale segmentation predictions than nnFormer, UNETR, and Swin UNETR in the illustrated cases.nnFormer confuses spleen with stomach, while UNETR and Swin UNETR confuse stomach with the background.
  • UNETR++ accurately segments all organs, whereas other methods under-segment the left adrenal gland or spleen and UNETR over-segments the stomach.

B.2. ACDC Dataset

On the ACDC dataset, qualitative comparisons show that UNETR++ more accurately segments the three heart categories than UNETR and nnFormer in the illustrated cases, including a difficult sample with smaller segment sizes.

  • Qualitative comparison: UNETR++ accurately segments all three categories in the first illustrated case, whereas UNETR and nnFormer under-segment the right ventricular cavity.The inaccurate predictions are marked with red dashed boxes.
  • Qualitative comparison: In a difficult sample where all three heart segments are comparatively smaller, Fig. 6 compares UNETR++ with nnFormer and UNETR.

B.3. Detailed qualitative comparison between UNETR++ and the baseline

Qualitative comparisons across Synapse, ACDC, and Decathlon-Lung show that UNETR++ more accurately delineates organs and cardiac regions than UNETR and other existing methods. The comparisons particularly highlight reduced under-segmentation, improved boundary delineation, and fewer false positives.

  • Synapse qualitative comparison: UNETR++ more precisely delineates the pancreas, veins, aorta, stomach, spleen, and left kidney than the UNETR baseline.UNETR under-segments or poorly delineates these structures, while UNETR++ segments them precisely with improved boundaries.
  • Synapse qualitative comparison: UNETR++ accurately segments organs on Synapse, whereas UNETR, Swin UNETR, and nnFormer struggle with different organ regions.The comparison marks inaccurate regions with red dashed boxes and reports promising segmentation performance for UNETR++.
  • ACDC qualitative comparison: UNETR++ segments all three ACDC right-ventricular-cavity regions more precisely than UNETR, which suffers from under-segmentation and poor boundary delineation.The qualitative comparison reports this pattern across all three rows.
  • Decathlon-Lung qualitative comparison: UNETR++ produces fewer false positives on Decathlon-Lung, while UNETR under-segments the entire region in the third row.The comparison is shown across multiple rows of the Decathlon-Lung examples.

C. Additional Ablations

Additional ablations show that UNETR++ scales to wider feature maps while improving BTCV segmentation, and that its EPA block outperforms alternative attention modules on Synapse. Qualitative results on Decathlon-Lung also show better tumor segmentation with fewer false positives than baseline UNETR.

  • Scalability: 84.27% average DSC is achieved with wider feature maps, improving from 83.28% while increasing parameters to 94.24M and FLOPs to 117G.The wider configuration uses feature-map sizes [64, 128, 256, 512] instead of [32, 64, 128, 256] on BTCV, without ensemble, pre-training, or additional custom data.
  • Qualitative comparison: UNETR++ provides better tumor segmentation and fewer false positives than baseline UNETR on Decathlon-Lung.The qualitative comparison highlights these differences in the enlarged tumor regions.
  • EPA ablation: UNETR++ achieves superior results on Synapse compared with gated attention and squeeze-and-excitation alternatives.The experiments replace or integrate EPA with gated attention from attention-gated U-Net and squeeze-and-excitation modules.

D. Discussion

UNETR++ is presented as a hierarchical segmentation approach that achieves promising results across five datasets while reducing model complexity and memory consumption and improving inference speed. Its efficient paired attention block encodes inter-dependent spatial and channel features through spatial and channel attention.

  • D. Discussion: UNETR++ achieves promising segmentation results on five datasets: Synapse, ACDC, BTCV, BRaTs, and Decathlon-Lung.The discussion identifies these benchmarks as evidence of the approach’s segmentation performance.
  • D. Discussion: UNETR++ reduces model complexity and memory consumption compared with existing methods.The passage presents these reductions as advantages of the hierarchical design.
  • D. Discussion: UNETR++ improves inference speed compared with existing methods through its efficient hierarchical architecture.The passage explicitly links the approach with improved inference speed, without giving a numerical gain.
  • D. Discussion: The efficient paired attention block encodes enriched inter-dependent spatial and channel features using spatial and channel attention.EPA is described as the mechanism that jointly models spatial and channel information.
Loading 2212.04497v3…