Source-linked AI summary
U-Net Transformer: Self and Cross Attention for Medical Image Segmentation
Olivier Petit, Nicolas Thome, Clément Rambour, Luc Soler
TL;DR
Complex, low-contrast anatomical structures challenge U-Nets because limited receptive fields do not capture sufficient contextual information. U-Transformer combines a U-shaped architecture with self-attention for global encoder interactions and cross-attention for filtering skip features, outperforming U-Net and Attention U-Net on two abdominal CT datasets. The results support using both attention mechanisms for abdominal organ segmentation, especially difficult small and complex organs.
Problem
U-Nets and related FCNs struggle with complex segmentation involving local ambiguities and low contrast because their receptive fields lack sufficient global context.
Method
U-Transformer augments a U-shaped FCN with self-attention for long-range interactions and cross-attention that filters non-semantic skip-connection features.
Results
U-Transformer outperforms U-Net by 2.4pts on TCIA and 1.3pts for IMO, and Attention U-Net by 1.7pts for TCIA and 1.6pts for IMO in Dice.
Takeaways & Limitations
The combined self- and cross-attention mechanisms are complementary, supporting better segmentation predictions, particularly for small and complex abdominal organs.
Abstract
from arXiv · showhide
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.
1 Introduction
U-Nets support medical organ segmentation but struggle with complex, low-contrast structures because limited receptive fields miss necessary context. U-Transformer addresses this gap by adding self- and cross-attention to model global anatomical relationships and refine decoder features.
- Organ segmentation supports treatment assessment and computer-assisted interventions in medical imaging.
- U-Nets use encoder-decoder processing, with convolutional encoders extracting semantic features and skip connections restoring high-resolution spatial information.
- Complex segmentation remains difficult for FCNs when local visual ambiguities and low contrast obscure anatomical boundaries.
- U-Transformer adds encoder self-attention for full contextual information and skip-connection cross-attention to filter non-semantic features during spatial recovery.
- Attention maps use long-range relationships with surrounding organs to support pancreas segmentation, while experiments compare U-Transformer with U-Net and local Attention U-Net methods.
2 The U-Transformer Network
U-Transformer augments a U-shaped encoder-decoder with multi-head self-attention and cross-attention. Self-attention supplies whole-image interactions, while cross-attention filters high-resolution skip features using higher-level semantics.
- U-Transformer uses Multi-Head Self-Attention and Multi-Head Cross-Attention to model long-range contextual interactions and spatial dependencies.
- Self-attention: MHSA sits at the encoder bottom and connects every element of the highest-level feature map, giving each pixel access to the whole input image.
- Self-attention: Self-attention embeds features into queries Q, keys K, and values V; similarities between Q and K weight value features to account for their interactions.
- Cross-attention: Skip connections preserve resolution but lack deep semantic richness, motivating MHCA to suppress irrelevant areas and emphasize application-relevant regions.
- Cross-attention: MHCA gates skip connection S using attention from high-level feature map Y, producing filter Z whose sigmoid-scaled weights yield Z ⊙ S before concatenation with Y.
3 Experiments
U-Transformer was evaluated for abdominal organ segmentation on TCIA pancreas and IMO CT datasets against U-Net and Attention U-Net. Results show gains on difficult organs and support complementary roles for self- and cross-attention, positional encoding, and multi-level cross-attention.
- 3.1 U-Transformer performances: U-Transformer outperformed U-Net by 2.4pts on TCIA and 1.3pts on IMO, and Attention U-Net by 1.7pts on TCIA and 1.6pts on IMO in Dice.Gains were consistent across folds, with paired t-tests showing p-values < 3% for every experiment.
- 3.1 U-Transformer performances: U-Transformer performed better on qualitative difficult cases with ambiguous local structures, including a complex pancreas shape missed by both baselines.Figure 5 compares U-Net, Attention U-Net, and U-Transformer on IMO and TCIA pancreas results.
- 3.1 U-Transformer performances: Self-attention and cross-attention alone outperformed U-Net and Attention U-Net on both datasets, while their combination produced complementary improvements.The comparison supports modeling global anatomical interactions and positional information rather than relying only on local skip-connection attention.
- 3.1 U-Transformer performances: U-Transformer significantly outperformed both baselines on challenging IMO organs: pancreas by +3.4pts, gallbladder by +1.3pts, and stomach by +2.2pts.It was better for every organ, including liver, where U-Net already scored > 95%.
- 3.2 U-Transformer analysis and properties: Positional encoding improved MHSA by +0.7pt on TCIA and +0.6pt on IMO, while single-level MHCA improved by +1.7pts and +0.6pt, respectively.These ablations were conducted on one fold of each dataset.
- 3.2 U-Transformer analysis and properties: Multi-level MHCA improved results over single-level attention by +1.8pts on TCIA and +0.6pt on IMO.Higher-resolution level-3 attention maps tended to focus on more specific regions than earlier levels.
- 3.2 U-Transformer analysis and properties: More MHSA heads improved performance, with the largest gain from the first head, and Hausdorff distance followed the same trend as Dice.The authors relate this trend to reduced prediction artefacts through self- and cross-attention.
4 Conclusion
The paper introduces U-Transformer, which augments a U-shaped FCN with self- and cross-attention to model long-range interactions and spatial dependencies. The approach is highlighted for abdominal organ segmentation, especially for small and complex organs.
- 4 Conclusion: U-Transformer augments a U-shaped FCN with self- and cross-attention modules to model long-range interactions and spatial dependencies.The authors identify 3D networks, MRI or ultrasound, and other medical image tasks as future study areas.