Source-linked AI summary
UTNet: A Hybrid Transformer Architecture for Medical Image Segmentation
Yunhe Gao, Mu Zhou, Dimitris Metaxas
TL;DR
Medical vision segmentation needs long-range context, but standard Transformers incur quadratic attention costs and require substantial data to learn visual inductive bias. UTNet combines convolutional layers with efficient self-attention and relative position encoding across encoder and decoder stages. It reports superior segmentation performance and robustness on a multi-label, multi-vendor cardiac MRI cohort.
Problem
Medical vision applications of Transformers face quadratic self-attention overhead and limited medical data for learning visual inductive bias.
Method
UTNet integrates convolutional feature extraction with efficient self-attention and relative position encoding in a U-shaped encoder-decoder network.
Results
UTNet demonstrates superior segmentation performance and robustness on a multi-label, multi-vendor cardiac MRI cohort.
Takeaways & Limitations
The hybrid design allows Transformer initialization within convolutional networks without pre-training and supports efficient long-range dependency modeling.
Abstract
from arXiv · showhide
Transformer architecture has emerged to be successful in a number of natural language processing tasks. However, its applications to medical vision remain largely unexplored. In this study, we present UTNet, a simple yet powerful hybrid Transformer architecture that integrates self-attention into a convolutional neural network for enhancing medical image segmentation. UTNet applies self-attention modules in both encoder and decoder for capturing long-range dependency at different scales with minimal overhead. To this end, we propose an efficient self-attention mechanism along with relative position encoding that reduces the complexity of self-attention operation significantly from $O(n^2)$ to approximate $O(n)$. A new self-attention decoder is also proposed to recover fine-grained details from the skipped connections in the encoder. Our approach addresses the dilemma that Transformer requires huge amounts of data to learn vision inductive bias. Our hybrid layer design allows the initialization of Transformer into convolutional networks without a need of pre-training. We have evaluated UTNet on the multi-label, multi-vendor cardiac magnetic resonance imaging cohort. UTNet demonstrates superior segmentation performance and robustness against the state-of-the-art approaches, holding the promise to generalize well on other medical image segmentations.
1 Department of Computer Science, Rutgers University
The listed affiliation includes the Department of Computer Science at Rutgers University and identifies SenseBrain, Shanghai AI Laboratory, and the Centre for Perceptual and Interactive Intelligence.
- It also lists SenseBrain among the affiliations.
- Shanghai AI Laboratory and the Centre for Perceptual and Interactive Intelligence are also listed.
1 Introduction
The introduction motivates combining convolutional locality with Transformer self-attention to address limited global context, computational cost, and data demands in medical vision.
- Convolutional networks capture local texture features but lack explicit long-range dependency modeling.
- Self-attention captures long-range dependency and dynamically aggregates relevant features from image content.
- Standard self-attention has O(n^2) time and space complexity, restricting its use on high-resolution feature maps.
- UTNet combines convolutional feature extraction with efficient self-attention and relative position encoding for medical image segmentation.Its efficient attention reduces complexity from O(n^2) to approximate O(n), and evaluation covers multi-label, multi-vendor cardiac MRI.
2 Method
The method integrates multi-head self-attention into a convolutional U-shaped network, reducing attention cost while retaining global and positional information across resolutions.
- Self-attention: Multi-head self-attention projects feature maps into query, key, and value embeddings before global context aggregation.The spatial sequence length is n = HW, and the model uses four heads.
- Network architecture: The proposed decoder attention combines high-resolution skip-connection features with low-resolution decoder features.
- Efficient self-attention: The efficient mechanism projects keys and values from n positions into k ≪ n positions, reducing complexity to O(nkd).Down-sampling can use pooling or strided convolution; the implementation uses 1×1 convolution followed by bilinear interpolation.
- Relative positional encoding: Relative position encoding adds learnable relative height and width information to attention logits.The encoding is computed after low-dimensional projection and models content-position relationships.
- Network architecture: UTNet places self-attention across encoder and decoder resolutions while using convolutional inductive bias and pre-activation residual shortcuts.This design targets efficient processing of high-resolution feature maps and avoids large-scale pretraining.
3 Experiment
UTNet is evaluated against convolutional and attention-based segmentation models on cardiac MRI, including efficiency, ablation, and cross-vendor robustness. It leads segmentation outcomes while reducing memory and inference time relative to quadratic dual attention and retaining performance on unseen vendors.
- Segmentation performance: UTNet demonstrates leading performance for LV, MYO, and RV segmentation outcomes.
- Ablation study: Self-attention at progressively higher network levels improves performance by gathering more fine-grained detail, but the curve saturates at the original resolution.The x-axis notation indicates the down-sampling levels where self-attention is placed.
- Ablation study: A reduced size of 8 gives the best efficient self-attention performance, while interpolation down-sampling is slightly better than max-pooling.
- Ablation study: The Transformer encoder-decoder combination gives optimal performance, and removing relative positional encoding causes a large performance drop.
- Efficiency comparison: 3.8 GB vs 36.9 GB GPU memory and 0.146 s vs 0.243 s inference time favor UTNet over dual attention.The comparison uses four resolutions, identical input size of 256 × 256 × 16, and identical batch size.
- Robustness analysis: UTNet retains competitive performance on unseen vendors C and D, whereas competing approaches suffer from vendor differences.The evaluation trains on vendors A and B and tests on vendors A, B, C, and D.
4 Conclusion
UTNet combines convolutional layers with self-attention in a U-shaped hybrid network for medical image segmentation. Its design supports training from scratch, multi-level long-range dependency capture, and efficient handling of long sequences.
- UTNet merges convolutional layers and self-attention in a U-shaped hybrid Transformer network for medical image segmentation.
- The hybrid layer design allows Transformer initialization within convolutional networks without pre-training.
- Self-attention operates at multiple network levels in both encoder and decoder to capture long-range dependencies.
- Efficient long-sequence handling opens possibilities for applying UTNet to additional downstream medical image tasks.