PaperlayerView exact revision on arXiv

Source-linked AI summary

CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation

Yutong Xie, Jianpeng Zhang, Chunhua Shen, Yong Xia

arXiv:2103.03024v1cs.CV

TL;DR

CNNs are limited in modeling long-range dependencies, whereas vanilla Transformers are computationally and spatially expensive on high-resolution 3D feature maps. CoTr bridges them with CNN feature extraction and a deformable Transformer that attends to a small set of key positions. On the BCV 3D multi-organ segmentation task, CoTr outperforms CNN-based, Transformer-based, and hybrid methods while processing multi-scale high-resolution features.

  • Problem

    CNNs have limited long-range dependency modeling, while vanilla Transformers are too computationally and spatially complex for high-resolution 3D feature maps.

  • Method

    CoTr combines a CNN encoder with a deformable Transformer using deformable self-attention to model long-range dependencies on multi-scale feature maps.

  • Results

    CoTr outperforms CNN-based, vanilla Transformer-based, and hybrid models on 3D multi-organ segmentation, improving average Dice by 1.4% over CoTr without DeTrans.

  • Takeaways & Limitations

    Deformable self-attention enables CoTr to combine CNN low-level details with Transformer long-range modeling using multi-scale features.

Abstract

from arXiv · show

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{\rm\small\ttfamily} \url{https://github.com/YtongXie/CoTr}

1 Introduction

3D medical image segmentation is dominated by CNNs, but their locality limits long-range modeling, while vanilla Transformers are too costly for high-resolution 3D features. CoTr combines CNN feature extraction with an efficient deformable Transformer that processes multi-scale features.

  • CNNs are prevalent in medical image segmentation but have limited receptive fields for long-range dependency modeling.
  • Transformers model long-range dependencies through content-adaptive self-attention but incur high training, computational, and spatial costs in 3D.
  • Hybrid CNN-Transformer models combine CNN feature extraction with Transformer-based contextual modeling for medical image segmentation.
  • CoTr uses deformable self-attention over a small set of key sampling points to reduce Transformer complexity.
  • CoTr processes multi-scale and high-resolution feature maps while retaining long-range dependency modeling for 3D segmentation.

2 Materials

The study uses the BCV dataset for automated segmentation of 11 abdominal organs from 30 labeled CT scans.

  • The BCV dataset contains 30 labeled CT scans for automated abdominal-organ segmentation.
  • The dataset covers 11 organs, including the spleen, kidneys, gallbladder, liver, stomach, vessels, pancreas, and adrenal gland.
  • The evaluated anatomy includes abdominal organs and vascular structures such as the aorta, inferior vena cava, and portal and splenic veins.

3 Methods

CoTr combines a CNN encoder for feature extraction with a deformable Transformer encoder for long-range dependency modeling, followed by a CNN decoder for segmentation. Multi-scale deformable attention, positional encoding, skip connections, and deep supervision support efficient processing and reconstruction of 3D feature maps.

  • 3.1 CNN-encoder: CoTr uses a CNN encoder, DeTrans-encoder, and decoder to learn representations and produce 3D segmentations.The DeTrans-encoder models long-range dependencies between CNN features.
  • 3.1 CNN-encoder: The CNN encoder contains a Conv-IN-ReLU block and three stages of 3D residual blocks with 3, 3, and 2 blocks.The feature maps are represented across L levels with C channels and spatial dimensions reduced by powers of two.
  • 3.2 DeTrans-encoder: Because convolution is intrinsically local, the DeTrans-encoder adds multi-scale deformable self-attention for efficient long-range contextual modeling.It flattens multi-scale CNN features into a sequence and supplements them with 3D positional encodings.
  • 3.2 DeTrans-encoder: 3D positional encodings use sine and cosine functions across depth, height, and width, then are added element-wise to flattened features.The encoding combines PED, PEH, and PEW for each feature level.
  • 3.2 DeTrans-encoder: MS-DMSA computes each attention head from multi-scale features using K sampled key points, learned attention weights, offsets, and rescaled reference coordinates.The attention weights and sampling offsets are obtained by linear projection of the query feature.
  • 3.2 DeTrans-encoder: Each DeTrans layer combines MS-DMSA with a feed-forward network, layer normalization, and skip connections, and layers are repeatedly stacked.The attention heads are linearly projected to aggregate their representations.
  • 3.3 Decoder: The CNN decoder reshapes DeTrans outputs, progressively upsamples them to D × H × W, refines them with residual blocks, and uses skip connections.Deep supervision adds auxiliary losses at different decoder scales, while training minimizes Dice plus cross-entropy loss.
  • 3.4 Implementation: Training uses cropped 48×192×192 CT sub-volumes and augmentation, while testing uses sliding windows, Gaussian weighting, and axis-flip test-time augmentation.The BCV split contains 21 training scans, 9 test scans, and 6 training scans selected for validation.

4 Results

CoTr is evaluated against pure Transformer, pure CNN, and hybrid CNN-Transformer models on BCV, with additional analysis of efficiency and multi-scale feature use.

  • Comparing to models with only Transformer encoder: CoTr improves average Dice by 4% over the comparable-parameter CoTr w/o CNN-encoder variant.The comparison uses the same decoder and evaluates segmentation across 11 organs.
  • Comparing to models with only CNN encoder: CoTr improves average Dice by 1.4% over CoTr w/o DeTrans and consistently outperforms the evaluated CNN-based context-modeling alternatives.The alternatives replace DeTrans with ASPP, pyramid parsing, or Non-local modules while retaining the same CNN-encoder and decoder.
  • Hyper-parameter and feature-scale analysis: Figure 3 reports validation-set average Dice against sampled key points, heads, DeTrans layers, and single-scale versus multi-scale features.The figure organizes these comparisons across panels (a)–(d).
  • Comparing to models with hybrid CNN-Transformer encoder: CoTr steadily outperforms TransUNet across all organs, while the smaller CoTr† variant performs better on seven organs.The reported advantage is attributed to deformable attention reducing computational and spatial complexity for high-resolution, multi-scale feature maps.
  • Computational Complexity: CoTr takes about 2 days to train and less than 30ms to segment a 48 × 192 × 192 volume.The measurements were obtained using an NVIDIA GTX 2080Ti GPU and PyTorch.

5 Discussion on Hyper-parameter Settings

The study examines how DeTrans hyper-parameters and feature-scale choices affect validation performance.

  • Hyper-parameter settings: Increasing sampled key points K, heads H, or DeTrans layers LD improves validation-set average Dice over all organs.The tested settings are K = 1, 2, 4; H = 2, 4, 6; and LD = 2, 4, 6.
  • Feature-scale settings: Multi-scale feature maps improve average Dice by 1.2% compared with single-scale feature maps from the last stage.The comparison uses CoTr trained with either single-scale or multi-scale features.

6 Conclusion

CoTr combines CNN feature extraction with a deformable Transformer to model long-range dependencies for 3D medical image segmentation.

  • Conclusion: CoTr uses a hybrid CNN-Transformer encoder that balances preserving low-level details with modeling long-range dependency.Its deformable mechanism reduces computational and spatial complexities, enabling high-resolution and multi-scale feature processing.
  • Conclusion: Comparative BCV experiments show superior performance over CNN-based and vanilla Transformer-based models.The paper identifies CoTr as a stronger baseline that could be extended to brain or tumor segmentation.

7 Appendix

The appendix details CoTr’s CNN encoder, decoder, feed-forward components, training loss, and visual segmentation comparisons. It specifies the architectural building blocks and reports that CoTr produces results closer to ground truth with fewer false positives than compared methods.

  • Detailed network architecture: The encoder uses a Conv-In-ReLU block and three stages of 3D residual blocks, while the decoder uses transposed convolutions, residual refinement, and encoder skip connections.The residual-block stages contain three, three, and two blocks; the decoder’s first three upsampling modules combine TransConv, residual blocks, and pixel-wise summation.
  • Detailed network architecture: Blue Conv blocks contain 3D convolution, instance normalization, and ReLU, whereas gray Conv and orange TransConv denote 3D convolutional and transposed-convolutional layers.Numbers in Conv blocks or layers indicate filters, kernel size, and stride; residual-block numbers indicate filters, stride, and Conv-block count.
  • Loss function: The optimization combines soft Dice loss with cross-entropy loss and applies deep supervision by summing auxiliary losses from decoder outputs at different resolutions.The prediction and ground truth are denoted by ỹ and y, respectively, while ε is a smoothing factor and c is the number of categories.
  • Visualization: CoTr’s visual results are more similar to ground truth than SETR and ASPP, and more likely to contain fewer false positives than TransUNet.Figure 5 compares four cases and marks regions indicating CoTr’s superiority with yellow rectangles; organ types use distinct colors.
Loading 2103.03024v1…