[Paper Review] CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation

Yutong Xie, Jianpeng Zhang+2 more|arXiv (Cornell University)|Mar 4, 2021

Advanced Neural Network ApplicationsComputer Science28 references56 citations

TL;DR

CoTr hybrids a CNN encoder with a deformable Transformer to efficiently model long-range context for 3D medical image segmentation, achieving state-of-the-art performance on BCV 11-organ segmentation while handling high-resolution multi-scale features.

ABSTRACT

Convolutional neural networks (CNNs) have been the de facto standard for nowadays 3D medical image segmentation. The convolutional operations used in these networks, however, inevitably have limitations in modeling the long-range dependency due to their inductive bias of locality and weight sharing. Although Transformer was born to address this issue, it suffers from extreme computational and spatial complexities in processing high-resolution 3D feature maps. In this paper, we propose a novel framework that efficiently bridges a {\bf Co}nvolutional neural network and a {\bf Tr}ansformer {\bf (CoTr)} for accurate 3D medical image segmentation. Under this framework, the CNN is constructed to extract feature representations and an efficient deformable Transformer (DeTrans) is built to model the long-range dependency on the extracted feature maps. Different from the vanilla Transformer which treats all image positions equally, our DeTrans pays attention only to a small set of key positions by introducing the deformable self-attention mechanism. Thus, the computational and spatial complexities of DeTrans have been greatly reduced, making it possible to process the multi-scale and high-resolution feature maps, which are usually of paramount importance for image segmentation. We conduct an extensive evaluation on the Multi-Atlas Labeling Beyond the Cranial Vault (BCV) dataset that covers 11 major human organs. The results indicate that our CoTr leads to a substantial performance improvement over other CNN-based, transformer-based, and hybrid methods on the 3D multi-organ segmentation task. Code is available at \def\UrlFont{ m\small tfamily} \url{https://github.com/YtongXie/CoTr}

Motivation & Objective

• Motivate bridging CNNs’ local inductive bias with Transformers for 3D medical image segmentation.
• Develop an efficient deformable Transformer (DeTrans) to model long-range dependencies on multi-scale feature maps.
• Design a CNN-encoder–DeTrans-encoder–decoder architecture to preserve high-resolution details while capturing global context.
• Demonstrate improved segmentation performance over CNN-based, Transformer-based, and other hybrid methods on BCV dataset.

Proposed method

• Use a CNN-encoder to extract multi-scale 3D feature maps.
• Introduce a DeTrans-encoder with multi-scale deformable self-attention to efficiently capture long-range dependencies.
• Flatten CNN features with 3D positional encoding and process via DeTrans layers.
• Apply a multi-head deformable self-attention with limited sampling points to reduce complexity.
• Fuse DeTrans outputs through a CNN-based decoder with skip connections and deep supervision.
• Optimize with a joint Dice and cross-entropy loss; employ data augmentation and instance normalization.

Experimental results

Research questions

• RQ1 Can a lightweight hybrid CNN–Transformer encoder, augmented with deformable self-attention, outperform pure CNN or pure Transformer approaches in 3D medical image segmentation?
• RQ2 Does multi-scale deformable self-attention enable effective long-range modeling on high-resolution 3D feature maps?
• RQ3 What is the impact of DeTrans hyper-parameters and multi-scale feature integration on segmentation performance?
• RQ4 How does CoTr compare to existing CNN-, Transformer-, and hybrid-based methods on the BCV multi-organ segmentation task?

Key findings

• CoTr outperforms CNN-only, Transformer-only, and other hybrid baselines on BCV 3D multi-organ segmentation.
• CoTr with deformable self-attention enables processing of multi-scale high-resolution feature maps with reduced computational and spatial complexity.
• CoTr variants with smaller CNN encoders (CoTr ∗, CoTr †) achieve strong results, showing the benefit of the hybrid encoder over purely Transformer-based encoders.
• Replacing DeTrans with traditional context modules (ASPP, PP, Non-local) leads to lower Dice scores, highlighting the advantage of deformable Transformers.
• CoTr consistently improves average Dice across 11 organs, notably for gallbladder and pancreas, and achieves competitive or superior performance to TransUNet in 3D settings.
• Training-time and inference-time efficiency: training ~2 days on GTX 2080Ti; inference under 30 ms per 48×192×192 volume.