Source-linked AI summary
TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
Jieneng Chen, Yongyi Lu, Qihang Yu, Xiangde Luo, Ehsan Adeli, Yan Wang, Le Lu, Alan L. Yuille, Yuyin Zhou
TL;DR
Medical image segmentation needs both long-range context and precise localization, which standalone CNNs and Transformers each handle incompletely. TransUNet combines a Transformer encoder with a U-Net-style decoder and CNN skip features, achieving superior performance across various medical image segmentation tasks.
Problem
Medical image segmentation requires long-range dependency modeling and fine localization, while CNNs and Transformers each have limitations in one of these capabilities.
Method
TransUNet combines a Transformer encoder for global context with a hybrid CNN-U-Net decoder using high-resolution skip features for precise localization.
Results
TransUNet achieves superior performance to various competing methods across diverse medical image segmentation tasks, including multi-organ and cardiac segmentation.
Takeaways & Limitations
The study supports Transformers as strong encoders for medical image segmentation when paired with U-Net-style recovery of localized spatial information.
Abstract
from arXiv · showhide
Medical image segmentation is an essential prerequisite for developing healthcare systems, especially for disease diagnosis and treatment planning. On various medical image segmentation tasks, the u-shaped architecture, also known as U-Net, has become the de-facto standard and achieved tremendous success. However, due to the intrinsic locality of convolution operations, U-Net generally demonstrates limitations in explicitly modeling long-range dependency. Transformers, designed for sequence-to-sequence prediction, have emerged as alternative architectures with innate global self-attention mechanisms, but can result in limited localization abilities due to insufficient low-level details. In this paper, we propose TransUNet, which merits both Transformers and U-Net, as a strong alternative for medical image segmentation. On one hand, the Transformer encodes tokenized image patches from a convolution neural network (CNN) feature map as the input sequence for extracting global contexts. On the other hand, the decoder upsamples the encoded features which are then combined with the high-resolution CNN feature maps to enable precise localization. We argue that Transformers can serve as strong encoders for medical image segmentation tasks, with the combination of U-Net to enhance finer details by recovering localized spatial information. TransUNet achieves superior performances to various competing methods on different medical applications including multi-organ segmentation and cardiac segmentation. Code and models are available at https://github.com/Beckschen/TransUNet.
1 Introduction
The introduction motivates TransUNet by contrasting U-Net’s strong detail retention with CNNs’ limited explicit long-range modeling and Transformers’ coarse localization. TransUNet addresses these limitations through a hybrid CNN-Transformer architecture combining high-resolution spatial information with global context, and outperforms competing methods across medical segmentation tasks.
- Motivation: U-Net, a symmetric encoder-decoder with skip-connections, has become the de-facto choice for medical image segmentation.CNN- and FCN-based approaches have achieved success in cardiac, organ, and polyp segmentation.
- Motivation: CNNs generally struggle to model explicit long-range relations because convolution operations are intrinsically local.This limitation particularly affects structures with substantial inter-patient variation in texture, shape, and size.
- Motivation: A naive Transformer encoder followed by direct upsampling produces unsatisfactory, coarse segmentation because low-resolution global features lack recoverable localization details.Transformers treat inputs as 1D sequences and focus on global context at all stages, while direct upsampling cannot effectively restore detailed localization information.
- TransUNet: TransUNet uses a hybrid CNN-Transformer architecture to combine detailed high-resolution CNN spatial information with global context encoded by Transformers.The framework establishes self-attention mechanisms from a sequence-to-sequence prediction perspective and compensates for Transformer-induced feature-resolution loss.
- Results: TransUNet achieves superior performance against competing methods across various medical image segmentation tasks.The introduction reports extensive experimental evidence supporting the method’s superiority on different medical applications.
2 Related Works
Prior work integrates self-attention into CNN-based segmentation architectures and adapts Transformers for computer vision by restricting or approximating global attention. These approaches motivate combining convolutional representations with mechanisms for modeling broader interactions.
- Combining CNNs with self-attention mechanisms: Prior studies integrated self-attention into CNNs by modeling global pixel interactions and adding attention modules to encoder-decoder architectures.Examples include non-local operators inserted into intermediate convolution layers and additive attention gates integrated into skip connections.
- Transformers: Transformers, originally proposed for machine translation, were adapted for vision through local-neighborhood attention and scalable approximations to global self-attention.The passage cites local attention for each query pixel and Sparse Transformers as examples of these modifications.
3 Method
TransUNet combines a CNN-Transformer hybrid encoder with a cascaded upsampler to model global context while recovering high-resolution spatial details for precise segmentation. The encoder tokenizes CNN feature maps, and the decoder aggregates multiscale features through skip-connections.
- Transformer Encoder: TransUNet introduces self-attention into the encoder by applying Transformers to tokenized image representations.The Transformer encoder uses stacked Multihead Self-Attention and Multi-Layer Perceptron blocks to produce an encoded image representation.
- Decoder Design: Directly bilinearly upsampling the encoded representation is a naive decoder baseline denoted “None” in later comparisons.This approach reshapes the encoded feature, reduces its channels to the number of classes, and upsamples it to H × W.
- Cascaded Upsampler: The hybrid encoder and cascaded upsampler form a U-shaped architecture that aggregates features at different resolutions through skip-connections.These connections recover localized spatial information lost when Transformer features originate from lower-resolution patches.
- CNN-Transformer Hybrid as Encoder: The hybrid encoder first extracts CNN feature maps, then applies patch embedding to 1 × 1 patches from those maps rather than raw images.This design enables use of intermediate high-resolution CNN feature maps during decoding and outperforms a pure Transformer encoder.
- Cascaded Upsampler: The cascaded upsampler decodes hidden features through multiple 2× upsampling, 3×3 convolution, and ReLU blocks to reach full resolution.It progressively maps features from H/P to H×W.
4 Experiments and Discussion
Experiments on Synapse multi-organ CT and ACDC cardiac MRI evaluate TransUNet against established segmentation methods and isolate the contributions of its decoder, hybrid encoder, skip-connections, resolution, patch size, and model scale. TransUNet improves segmentation through global contextual modeling combined with recovered spatial detail, with higher-resolution inputs and larger models providing further gains at increased cost.
- Datasets and evaluation: Experiments use Synapse CT scans with 8 abdominal organs and ACDC cardiac MRI scans annotated for LV, RV, and MYO.Synapse contains 30 scans and 3779 axial images, while ACDC uses 70 training, 10 validation, and 20 testing cases.
- Main results: 6.36% average DSC and 3.50 mm Hausdorff-distance improvements from ViT-None to ViT-CUP show that CUP outperforms direct upsampling.R50-ViT-CUP adds another 3.43% DSC and 3.24 mm Hausdorff-distance improvement over ViT-CUP, demonstrating the hybrid encoder’s effectiveness.
- Main results: 1.91% to 8.67% average DSC gains over prior methods establish TransUNet’s superior Synapse performance, while ViT-CUP reaches 67.86% DSC but trails U-Net and AttnUNet.The results attribute the advantage to combining Transformer high-level semantics with low-level spatial cues.
- Ablation studies: 1.4% DSC improves performance when lightweight additive Transformers are inserted into a skip-connection, supporting enhanced recovery of fine segmentation details.The broader skip-connection ablation varies the number of connections among 0, 1, and 3.
- Ablation studies: 6.88% average DSC improves when input resolution increases from 224×224 to 512×512, but the larger sequence length substantially increases computational cost.With patch size 16, 512×512 inputs produce an approximate 5× larger Transformer sequence length; smaller patches and larger models also generally improve performance.
5 Conclusion
The paper presents TransUNet as an alternative framework for medical image segmentation that combines Transformers’ global self-attention with low-level CNN features in a u-shaped hybrid design.
- TransUNet investigates Transformers for general medical image segmentation as the paper’s first such study.
- The framework encodes image features as sequences to capture strong global context.
- Its u-shaped hybrid architecture incorporates low-level CNN features alongside Transformer representations.