[Paper Review] U-Net Transformer: Self and Cross Attention for Medical Image Segmentation

Olivier Petit, Nicolas Thome+2 more|arXiv (Cornell University)|Mar 10, 2021

Radiomics and Machine Learning in Medical Imaging18 citations

TL;DR

This paper proposes U-Transformer, a U-Net-based architecture enhanced with multi-head self-attention and cross-attention mechanisms to improve medical image segmentation, particularly for low-contrast and complex organs. By modeling long-range contextual dependencies and refining spatial details through attention-guided skip connections, U-Transformer achieves state-of-the-art performance on abdominal CT datasets, outperforming U-Net and Attention U-Net by up to 3.4% in Dice score on challenging organs like the pancreas.

ABSTRACT

Medical image segmentation remains particularly challenging for complex and low-contrast anatomical structures. In this paper, we introduce the U-Transformer network, which combines a U-shaped architecture for image segmentation with self- and cross-attention from Transformers. U-Transformer overcomes the inability of U-Nets to model long-range contextual interactions and spatial dependencies, which are arguably crucial for accurate segmentation in challenging contexts. To this end, attention mechanisms are incorporated at two main levels: a self-attention module leverages global interactions between encoder features, while cross-attention in the skip connections allows a fine spatial recovery in the U-Net decoder by filtering out non-semantic features. Experiments on two abdominal CT-image datasets show the large performance gain brought out by U-Transformer compared to U-Net and local Attention U-Nets. We also highlight the importance of using both self- and cross-attention, and the nice interpretability features brought out by U-Transformer.

Motivation & Objective

• To address the limitation of U-Net in capturing long-range contextual dependencies and spatial relationships in low-contrast, complex anatomical structures.
• To improve segmentation accuracy for small and visually ambiguous organs such as the pancreas, gallbladder, and stomach.
• To integrate both self-attention and cross-attention mechanisms into a U-Net framework to model global context and refine spatial details.
• To enhance model interpretability through attention map visualization and ablation studies.
• To validate the complementary benefits of self- and cross-attention in a unified transformer-augmented U-Net architecture.

Proposed method

• Introduces a U-Transformer network that combines the U-Net encoder-decoder architecture with multi-head self-attention (MHSA) at the end of the encoder to model global contextual interactions across the entire image.
• Employs multi-head cross-attention (MHCA) in the skip connections to fuse high-level semantic features with high-resolution features, enabling fine-grained spatial recovery.
• Uses learnable positional encoding to preserve spatial location information in both MHSA and MHCA modules, improving attention localization.
• Applies multi-level cross-attention by stacking MHCA modules at different decoder stages to progressively refine predictions.
• Utilizes a dual-branch attention mechanism where MHSA captures long-range dependencies between features, and MHCA filters non-semantic noise during feature up-sampling.
• Employs a standard U-Net backbone with residual blocks and skip connections, enhanced with attention modules instead of replacing convolutional layers.

Experimental results

Research questions

• RQ1 Can self-attention in the encoder improve segmentation of complex organs by modeling long-range contextual dependencies beyond the limited receptive field of U-Net?
• RQ2 Does cross-attention in the skip connections enhance spatial detail recovery by filtering out non-semantic features and integrating positional information?
• RQ3 How do self- and cross-attention mechanisms compare in isolation and in combination in terms of segmentation performance and interpretability?
• RQ4 What is the impact of positional encoding and multi-level attention on the performance of cross-attention in the decoder path?
• RQ5 Can the integration of both attention mechanisms lead to better generalization across organs, especially those with low contrast and complex shapes?

Key findings

• U-Transformer achieves a Dice similarity coefficient (DSC) of 73.10% on the pancreas in the IMO dataset, representing a +3.4% improvement over U-Net and a +4.45% improvement over Attention U-Net.
• On the TCIA pancreas dataset, U-Transformer achieves a DSC of 80.65%, outperforming U-Net (76.35%) and Attention U-Net (77.23%) by significant margins.
• The combination of self- and cross-attention in U-Transformer yields superior results compared to either mechanism alone, demonstrating their complementary nature.
• Positional encoding improves MHCA performance by +1.7 points on TCIA and +0.6 points on IMO, highlighting its importance for spatial localization.
• Multi-level cross-attention improves performance by +1.8 points on TCIA and +0.6 points on IMO, showing that hierarchical attention refinement enhances segmentation accuracy.
• The Hausdorff distance results follow the same trend as Dice scores, indicating that U-Transformer reduces prediction artifacts and improves boundary precision.