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

Introduction

Image segmentation is a longstanding challenge in medical image analysis. Since the introduction of U-Net , fully convolutional neural networks (CNNs) have become the predominant approach to addressing this task . Despite their prevalence, CNNs still suffer from the limited receptive field and fail to capture the long-range dependency, due to the inductive bias of locality and weight sharing . Many efforts have been devoted to enlarge a CNN’s receptive field thus improve its ability to context modeling. Yu et al. proposed the atrous convolution with an adjustable dilated rate, which shows superior performance in semantic segmentation . More straightforwardly, Peng et al. designed large kernels to capture rich global context information. Zhao et al. employed the pyramid pooling at multiple feature scales to aggregate multi-scale global information. Wang et al. presented the non-local operations which is usually embedded at the end of encoder to capture the long-range dependency. Although improving the context modeling to some extent, these models still have an inevitably limited receptive field, stranded by the CNN architecture.

Transformer, a sequence-to-sequence prediction framework, has a proven track record in machine translation and nature language processing , due to its strong ability to long-range modeling. The self-attention mechanism in Transformer can dynamically adjust the receptive field according to the input content, and hence is superior to convolutional operations in modeling the long-range dependency.

Recently, Transformer has been considered as an alternative architecture, and has achieved competitive performance on many computer vision tasks, like image recognition , semantic/instance segmentation , object detection , low-level vision , and image generation . A typical example is the vision Transformer (ViT) , which outperforms a ResNet-based CNN on recognition tasks but at a cost of using 300M data for training. Since a huge training dataset is not always available, recent studies attempt to combine a CNN and a Transformer into a hybrid model. Carion et al. employed a CNN to extract image features and a Transformer to further process the extracted features. Chen et al. designed TransUNet, in which a CNN and a Transformer are combined in a cascade manner to make a strong encoder for 2D medical image segmentation. Although the design of TransUNet is interesting and the performance is good, it is challenging to optimize this model due to the existence of self-attention . First, it requires extremely long training time to focus the attention, which was initially cast to each pixel uniformly, on salient locations, especially in a 3D scenario. Second, due to its high computational complexity, a vanilla Transformer can hardly process multi-scale and high-resolution feature maps, which play a critical role in image segmentation.

In this paper, we propose a hybrid framework that efficiently bridges Co-nvolutional neural network and Transformer (CoTr) for 3D medical image segmentation. CoTr has an encoder-decoder structure. In the encoder, a concise CNN structure is adopted to extract feature maps and a Transformer is used to capture the long-range dependency (see Fig. 1). Inspired by , we introduce the deformable self-attention mechanism to the Transformer. This attention mechanism casts attentions only to a small set of key sampling points, and thus dramatically reduces the computational and spatial complexity of Transformer. As a result, it is possible for the Transformer to process the multi-scale feature maps produced by the CNN and keep abundant high resolution information for segmentation. The main contributions of this paper are three-fold: (1) we are the first to explore Transformer for 3D medical image segmentation, particularly in a computationally and spatially efficient way; (2) we introduce the deformable self-attention mechanism to reduce the complexity of vanilla Transformer, and thus enable our CoTr to model the long-range dependency using multi-scale features; (3) our CoTr outperforms the competing CNN-based, Transformer-based, and hybrid methods on the 3D multi-organ segmentation task.

Materials

The Multi-Atlas Labeling Beyond the Cranial Vault (BCV) dataset https://www.synapse.org/#!Synapse:syn3193805/wiki/217789 was used for this study. It contains 30 labeled CT scans for automated segmentation of 11 abdominal organs, including the spleen (Sp), kidney (Ki), gallbladder (Gb), esophagus (Es), liver (Li), stomach (St), aorta (Ao), inferior vena cava (IVC), portal vein and splenic vein (PSV), pancreas (Pa), and adrenal gland (AG).

Methods

CoTr aims to learn more effective representations for medical image segmentation via bridging CNN and Transformer. As shown in Fig. 2, it consists of a CNN-encoder for feature extraction, a deformable Transformer-encoder (DeTrans-encoder) for long-range dependency modeling, and a decoder for segmentation. We now delve into the details of each module.

The CNN-encoder $\mathcal{F}^{CNN}(\cdot)$ contains a Conv-IN-ReLU block and three stages of 3D residual blocks. The Conv-IN-ReLU block contains a 3D convolutional layer followed by an instance normalization (IN) and Rectified Linear Unit (ReLU) activation. The numbers of 3D residual blocks in three stages are three, three, and two, respectively.

Given an input image $\bm{x}$ with a height of $H$, a width of $W$, and a depth ($i.e.$, number of slices) of $D$, the feature maps produced by $\mathcal{F}^{CNN}(\cdot)$ can be formally expressed as

where $L$ indicates the number of feature levels, $\bm{\Theta}$ denotes the parameters of the CNN-encoder, and $C$ denotes the number of channels.

2 DeTrans-encoder

Due to the intrinsic locality of convolution operations, the CNN-encoder cannot capture the long-range dependency of pixels effectively. To this end, we propose the DeTrans-encoder that introduces the multi-scale deformable self-attention (MS-DMSA) mechanism for efficient long-range contextual modeling. The DeTrans-encoder is a composition of an input-to-sequence layer and $L_{D}$ stacked deformable Transformer (DeTrans) layers.

Input-to-sequence Transformation. Considering that Transformer processes the information in a sequence-to-sequence manner, we first flatten the feature maps produced by the CNN-encoder $\{\bm{f}_{l}\}_{l=1}^{L}$ into a 1D sequence. Unfortunately, the operation of flattening the features leads to losing the spatial information that is critical for image segmentation. To address this issue, we supplement the 3D positional encoding sequence $\{\bm{p}_{l}\}_{l=1}^{L}$ to the flattened $\{\bm{f}_{l}\}_{l=1}^{L}$. For this study, we use sine and cosine functions with different frequencies to compute the positional coordinates of each dimension $pos$, shown as follows

where $\#\in\left\{D,H,W\right\}$ indicates each of three dimensions, $\upsilon=1/10000^{2k/\frac{C}{3}}$. For each feature level $l$, we concatenate $PE_{D}$, $PE_{H}$, and $PE_{W}$ as the 3D positional encoding $\bm{p}_{l}$ and combine it with the flattened $\bm{f}_{l}$ via element-wise summation to form the input sequence of DeTrans-encoder.

MS-DMSA Layer. In the architecture of Transformer, the self-attention layer would look over all possible locations in the feature map. It has the drawback of slow convergence and high computational complexity, and hence can hardly process multi-scale features. To remedy this, we design the MS-DMSA layer that focuses only on a small set of key sampling locations around a reference location, instead of all locations.

where $H$ is the number of attention heads, and $\Phi(\cdot)$ is a linear projection layer that weights and aggregates the feature representation of all attention heads.

DeTrans Layer. The DeTrans layer is composed of a MS-DMSA layer and a feed forward network, each being followed by the layer normalization (see Fig. 2). The skip connection strategy is employed in each sub-layer to avoid gradient vanishing. The DeTrans-encoder is constructed by repeatedly stacking DeTrans layers.

3 Decoder

The output sequence of DeTrans-encoder is reshaped into feature maps according to the size at each scale. The decoder, a pure CNN architecture, progressively upsamples the feature maps to the input resolution ($i.e.$, $D\times H\times W$) using the transpose convolution, and then refines the upsampled feature maps using a 3D residual block. Besides, the skip connections between encoder and decoder are also added to keep more low-level details for better segmentation. We also use the deep supervision strategy by adding auxiliary losses to the decoder outputs with different scales. The loss function of our model is the sum of the Dice loss and cross-entropy loss . More details on the network architecture gare in Appendix.

4 Implementation details

Following , we first truncated the HU values of each scan using the range of $$ to filter irrelevant regions, and then normalized truncated voxel values by subtracting 82.92 and dividing by 136.97. We randomly split the BCV dataset into two parts: 21 scans for training and 9 scans for test, and randomly selected 6 training scans to form a validation set, which just was used to select the hyper-parameters of CoTr. The final results on the test set are obtained by the model trained on all training scans.

In the training stage, we randomly cropped sub-volumes of size $48\times 192\times 192$ from CT scans as the input. To alleviate the over-fitting of limited training data, we employed the online data argumentation , including the random rotation, scaling, flipping, adding white Gaussian noise, Gaussian blurring, adjusting rightness and contrast, simulation of low resolution, and Gamma transformation, to diversify the training set. Due to the benefits of instance normalization , we adopted the micro-batch training strategy with a small batch size of 2. To weigh the balance between training time cost and performance reward, CoTr was trained for 1000 epochs and each epoch contains 250 iterations. We adopted the stochastic gradient descent algorithm with a momentum of 0.99 and an initial learning rate of 0.01 as the optimizer. We set the hidden size in MS-DMSA and feed forward network to 384 and 1536, respectively, and empirically set the hyper-parameters $L_{D}=6$, $H=6$, and $K=4$. Besides, we formed two variants of CoTr with small CNN-encoders, denoted as CoTr∗ and CoTr†. In CoTr∗, there is only one 3D residual block in each stage of CNN-encoder. In CoTr†, the number of 3D residual blocks in each stage of CNN-encoder is two.

In the test stage, we employed the sliding window strategy, where the window size equals to the training patch size. Besides, Gaussian importance weighting and test time augmentation by flipping along all axes were also utilized to improve the robustness of segmentation. To quantitatively evaluate the segmentation results, we calculated the Dice coefficient scores (Dice) metric that measures the overlapping between a prediction and its ground truth.

Results

Comparing to models with only Transformer encoder. We first evaluated our CoTr against two variants of the state-of-the-art SEgmentation Transformer (SETR) , which were formed by using randomly initialized and pre-trained ViT-B/16 as the encoder. We also compared to a variant of CoTr that removes the CNN-encoder (CoTr w/o CNN-encoder). To ensure an unprejudiced comparison, all models use the same decoder. The segmentation performance of these models is shown in Table 1, from which three conclusions can be drawn. First, although the Transformer architecture is not limited by the type of input images, the ViT-B/16 pre-trained on 2D natural images does not work well on 3D medical images. The suboptimal performance may be attributed to the domain shift between 2D natural images and 3D medical images. Second, ‘CoTr w/o CNN-encoder’ has about 22M parameters and outperforms the SETR with about 100M parameters. We believe that a lightweight Transformer may be more friendly for medical image segmentation tasks, where there is usually a small training dataset. Third, our CoTr∗ with comparable parameters significantly outperforms ‘CoTr w/o CNN-encoder’, improving the average Dice over 11 organs by 4%. It suggests that the hybrid CNN-Transformer encoder has distinct advantages over the pure Transformer encoder in medical image segmentation.

Comparing to models with only CNN encoder. Then, we compared CoTr against a variant of CoTr that removes the DeTrans-encoder (CoTr w/o DeTrans) and three CNN-based context modeling methods, $i.e.$, the Atrous Spatial Pyramid Pooling (ASPP) module, pyramid parsing (PP) module, and Non-local module. For a fair comparison, we used the same CNN-encoder and decoder but replaced our DeTrans-encoder with ASPP, PP, and Non-local modules, respectively. The results in Table 1 shows that our CoTr elevates consistently the segmentation performance over ‘CoTr w/o DeTrans’ on all organs and improves the average Dice by 1.4%. It corroborates that our CoTr using a hybrid CNN-Transformer encoder has a stronger ability than using a pure CNN encoder to learn effective representations for medical image segmentation. Moreover, comparing to these context modeling methods, our Transformer architecture contributes to more accurate segmentation.

Comparing to models with hybrid CNN-Transformer encoder. We also compared CoTr to other hybrid CNN-Transformer architectures like TransUNet . To process 3D images directly, we extended the original 2D TransUNet to a 3D version by using 3D CNN-encoder and decoder as done in CoTr. We also set the number of heads and layers of Transformer in 3D TransUNet to be the same as our CoTr. It shows in Table 1 that CoTr steadily beats TransUNet in the segmentation of all organs, particularly for the gallbladder and pancreas segmentation. Even with a smaller CNN-encoder, CoTr† still achieves better performance than TransUnet in the segmentation of seven organs. The superior performance owes to the deformable mechanism in CoTr that makes it possible to process high-resolution and multi-scale feature maps due to the reduced computational and spatial complexities.

Computational Complexity. The proposed CoTr was trained using a workstation with a NVIDIA GTX 2080Ti GPU and the Pytorch software packages. It took about 2 days for training, and less than 30ms to segment a volume of size $48\times 192\times 192$.

Discussion on Hyper-parameter Settings

In the DeTrans-encoder, there are three hyper-parameters, $i.e.$, $K$, $H$, and $L_{D}$, which represent the number of sampled key points, heads, and stacked DeTrans layers, respectively. To investigate the impact of their settings on the segmentation, we set $K$ to 1, 2, and 4, set $H$ to 2, 4, and 6, and set $L_{D}$ to 2, 4, and 6. In Fig. 3 (a-c), we plotted the average Dice over all organs obtained on the validation set versus the values of $K$, $H$, and $L_{D}$. It shows that increasing the number of $K$, $H$, or $L_{D}$ can improve the segmentation performance. To demonstrate the performance gain resulted from the multi-scale strategy, we also attempted to train CoTr with single-scale feature maps from the last stage. The results in Fig. 3 (d) show that using multi-scale feature maps instead of single-scale feature maps can effectively improve the average Dice by 1.2%.

Conclusion

In this paper, we propose a hybrid model of CNN Transformer, namely CoTr, for 3D medical image segmentation. In this model, we design the deformable Transformer (DeTrans) that employs the deformable self-attention mechanism to reduce the computational and spatial complexities of modelling the long-range dependency on multi-scale and high-resolution feature maps. Comparative experiments were conducted on the BCV dataset. The superior performance of our CoTr over both CNN-based and vanilla Transformer-based models suggests that, via combining the advantages of CNN and Transformer, the proposed CoTr achieves the balance in keeping the details of low-level features and modeling the long-range dependency. As a stronger baseline, our CoTr can be extended to deal with other structures (e.g., brain structure or tumor segmentation) in the future.

References

Appendix

Fig. 4 shows the architecture of CNN-encoder, decoder and feed forward network in Detrans-encoder. It consists of a Conv-In-Relu and three stages of 3D residual blocks. The numbers of 3D residual blocks are three, three, and two in three stages, respectively. The decoder contains four upsampling modules. Each of first three modules has a TransConv layer followed by a residual block, and a pixel-wise summation with the corresponding feature maps from the encoder and the TransConv layer. The last module comprises of an Upsampling layer followed by a 1 × 1 Conv layer that maps the 64-channel feature maps to the desired number of classes. The feed forward network in Detrans-encoder has two linear projection layers. The first layer is followed by a layer normalization layer and a Dropout layer. The second layer is followed by a Dropout layer.

2 Loss function

We jointly use the Dice loss and cross-entropy loss for optimization, which is popular in many medical image segmentation applications and has achieved prominent success . The loss function is formulated as

3 Visualization

The segmentation results produced by (1) SETR with pre-trained ViT-B/16, (2) replacing DeTrans-encoder with ASPP module, (3) 3D TransUNet, and (4) our CoTr, were visually compared in Fig. 5. We can see that: 1) comparing to the pure Transformer encoder method (SETR) and pure CNN encoder method (ASPP), our CoTr with the hybrid CNN-Transformer encoder is able to produce the segmentation results that are more similar to the ground truth, and 2) our CoTr are more likely to produce less false positives compared to TransUNet, which confirms the superiority of our 3D deformable Transformer over vanilla Transformer.