[Paper Review] TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

Jieneng Chen, Yongyi Lu+7 more|arXiv (Cornell University)|Feb 8, 2021

Advanced Neural Network Applications16 references3,794 citations

TL;DR

TransUNet combines CNN-based high-resolution features with Transformer global context to achieve state-of-the-art medical image segmentation, outperforming pure CNN and pure Transformer baselines across multiple datasets.

ABSTRACT

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.

Motivation & Objective

• Motivate why CNNs (U-Net) struggle with long-range dependencies in medical segmentation.
• Propose a hybrid CNN-Transformer encoder to leverage both high-resolution details and global context.
• Design a cascaded upsampling decoder with skip connections to recover fine spatial details.
• Demonstrate empirical gains over CNN-based and Transformer-based baselines on multiple medical imaging tasks.

Proposed method

• Tokenize image patches and encode with a Transformer to capture global context.
• Use a CNN feature map to supply high-resolution patches for Transformer embedding (hybrid encoder).
• Upsample Transformer features with a cascaded upsampler (CUP) and fuse via U-Net-like skip connections.
• Train with standard SGD on pretrained backbones; use 224x224 inputs with patch size 16 default.
• Compare “None” (naive upsampling) vs CUP decoder and different encoder choices.
• Provide ablations on skip connections, resolution, patch size, and model scale.

Experimental results

Research questions

• RQ1 Can Transformers serve as strong encoders for medical image segmentation when complemented with CNN-based fine details?
• RQ2 Does a hybrid CNN-Transformer encoder plus a cascaded upsampling decoder outperform pure Transformer or pure CNN baselines in medical segmentation tasks?
• RQ3 What is the impact of skip connections, input resolution, patch size, and model scale on segmentation quality?
• RQ4 How well does TransUNet generalize across CT multi-organ segmentation and cardiac MRI segmentation datasets?

Key findings

• TransUNet achieves state-of-the-art average Dice Score (DSC) on Synapse multi-organ CT with 77.48% when using a R50-ViT-CUP baseline and reaches 89.71 DSC on ACDC cardiac MRI dataset (per Table 5).
• Ablations show that adding skip connections at multiple CUP resolutions improves performance, with best results when skipping at 1/2, 1/4, and 1/8 scales.
• Hybrid encoder (CNN + ViT) outperforms pure ViT and pure CNN baselines, demonstrating the benefit of combining high-resolution CNN features with global Transformer context.
• CUP decoder significantly improves over naive upsampling, and larger model size yields better performance (Base vs Large in their tests).
• Higher input resolution (512x512) raises average DSC by about 6.88% at the cost of computation; patch size 16 (sequence length 196) performs better than larger patches.
• Qualitative results show TransUNet produces fewer false positives and preserves fine organ boundaries compared to CNN-only and other Transformer-based models.