CiT-Net: Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation

1 Introduction

Medical image segmentation refers to dividing a medical image into several specific regions with unique properties. Medical image segmentation results can not only achieve abnormal detection of human body regions but also be used to guide clinicians. Therefore, accurate medical image segmentation has become a key component of computer-aided diagnosis and treatment, patient condition analysis, image-guided surgery, tissue and organ reconstruction, and treatment planning. Compared with common RGB images, medical images usually suffer from the problems such as high density noise, low contrast, and blurred edges. So how to quickly and accurately segment specific human organs and lesions from medical images has always been a huge challenge in the field of smart medicine.

In recent years, with the rapid development of computer hardware resources, researchers have continuously developed many new automatic medical image segmentation algorithms based on a large number of experiments. The existing medical image segmentation algorithms can be divided into two categories: based on convolutional neural networks (CNNs) and based on the Transformer networks.

The early traditional medical image segmentation algorithms are based on manual features designed by medical experts using professional knowledge [Suetens, 2017]. These methods have a strong mathematical basis and theoretical support, but these algorithms have poor generalization for different organs or lesions of the human body. Later, inspired by the full convolutional networks (FCN) [Long et al., 2015] and the encoder-decoder, Ronnebreger et al. designed the UNet [Ronneberger et al., 2015] network that was first applied to medical image segmentation. After the network was proposed, its symmetric U-shaped encoder and decoder structure received widespread attention. At the same time, due to the small number of parameters and the good segmentation effect of the U-Net network, deep learning has made a breakthrough in medical image segmentation. Then a series of improved medical image segmentation networks are inspired based on the U-Net network, such as 2D U-Net++ [Zhou et al., 2018], ResDO-UNet [Liu et al., 2023], SGU-Net [Lei et al., March 2023], 2.5D RIU-Net [Lv et al., 2022], 3D Unet [C¸ ic¸ek et al., 2016], V-Net [Milletari et al., 2016], etc. Among them, Alom et al. designed R2U-Net [Alom et al., 2018] by combining UNet, ResNet [Song et al., 2020], and recurrent neural network (RCNN) [Girshick et al., 2014]. Then Gu et al. introduced dynamic convolution [Chen et al., 2020] into U-Net proposed CA-Net [Gu et al., 2020]. Based on U-Net, Yang et al. proposed DCU-Net [Yang et al., 2022] by referring to the idea of residual connection and deformable convolution [Dai et al., 2017].Lei et al. [Lei et al., 2022] proposed a network ASENet based on adversarial consistency learning and dynamic convolution.

The rapid development of CNNs in the field of medical image segmentation is largely due to the scale invariance and inductive bias of convolution operation. Although this fixed receptive field improves the computational efficiency of CNNs, it limits its ability to capture the relationship between distant pixels in medical images and lacks the ability to model medical images in a long range.

Aiming at the shortcomings of CNNs in obtaining global features of medical images, scholars have proposed a Transformer architecture. In 2017, Vaswani et al. [Vaswani et al., 2017] proposed the first Transformer network. Because of its unique structure, Transformer obtains the ability to process indefinite-length input, establish long-range dependency modeling, and capture global information. With the excellent performance of Transformer in NLP fields, ViT [Dosovitskiy et al., 2020] applied Transformer to the field of image processing for the first time. Then Chen et al. put forward TransUNet [Chen et al., 2021], which initiates a new period of Transformer in the field of medical image segmentation. Valanarasu et al. proposed MedT [Valanarasu et al., 2021] in combination with the gating mechanism. Cao et al. proposed a pure Transformer network Swin-Unet [Cao et al., 2021] for medical image segmentation, in combination with the shifted-window multi-head self-attention (SW-MSA) in Swin Transformer [Liu et al., 2021b]. Subsequently, Wang et al. designed the BAT [Wang et al., 2021a] network for dermoscopic images segmentation by combining the edge detection idea [Sun et al., 2022]. Hatamizadeh et al. proposed Swin UNETR [Tang et al., 2022] network for 3D brain tumor segmentation. Wang et al. proposed the UCTransNet [Wang et al., 2022] network that combines the channel attention with Transformer.

These methods can be roughly divided into based on the pure Transformer architecture and based on the hybrid architecture of CNNs and Transformer. The pure Transformer network realizes the long-range dependency modeling based on self-attention. However, due to the lack of inductive bias of the Transformer itself, Transformer cannot be widely used in small-scale datasets like medical images [Shamshad et al., 2022]. At the same time, Transformer architecture is prone to ignore detailed local features, which reduces the separability between the background and the foreground of small lesions or objects with large-scale changes in the medical image.

The hybrid architecture of CNNs and Transformer realizes the local and global information modeling of medical images by taking advantage of the complementary advantages of CNNs and Transformer, thus achieving a better medical image segmentation effect [Azad et al., 2022]. However, this hybrid architecture still suffers from the following two problems. First, it ignores the problems of organ deformation and lesion irregularities when modeling local features, resulting in weak local feature expression. Second, it ignores the correlation between the feature map space and the channels when modeling the global feature, resulting in inadequate expression of self-attention. To address the above problems, our main contributions are as follows:

• A novel dynamic deformable convolution (DDConv) is
• proposed. Through task adaptive learning, DDConv can flexibly change the weight coefficient and deformation offset of convolution itself. DDConv can overcome the problems of fixation of receptive fields and sharing of convolution kernel parameters, which are common problems of vanilla convolution and its variant convolutions, such as Atrous convolution and Involution, etc. Improves the ability to perceive tiny lesions and targets with large-scale changes in medical images.
• A new (shifted)-window adaptive complementary attention module ((S)W-ACAM) is proposed. (S)W-ACAM realizes the cross-dimensional global modeling of medical images through four parallel branches of weight coefficient adaptive learning. Compared with the current popular attention mechanisms, such as CBAM and NonLocal, (S)W-ACAM fully makes up for the deficiency of the conventional attention mechanism in modeling the cross-dimensional relationship between spatial and channels. It can capture the cross-dimensional longdistance correlation features in medical images, and enhance the separability between the segmented object and the background in medical images.
• A new parallel network structure based on dynamically adaptive CNNs and cross-dimensional feature fusion Transformer is proposed for medical image segmentation, called CiT-Net. Compared with the current popular hybrid architecture of CNNs and Transformer, CiT-Net can maximize the retention of local and global features in medical images. It is worth noting that CiT-Net not only abandons pre-training but also has fewer parameters and less computational costs, which are 11.58 M and 4.53 GFLOPs respectively.

Compared with the previous vanilla convolution [Ronneberger et al., 2015], dynamic convolution [Chen et al., 2020] [Li et al., 2021], and deformable convolution [Dai et al., 2017], our DDConv can not only adaptively change the weight coefficient and deformation offset of the convolution according to the medical image task, but also better adapt to the shape of organs and small lesions with large-scale changes in the medical image, and additionally, it can improve the local feature expression ability of the segmentation network. Compared with the self-attention mechanism in the existing Transformer architectures [Cao et al., 2021] [Wang et al., 2021a], our (S)W-ACAM requires fewer parameters and less computational costs while it's capable of capturing the global cross-dimensional long-range dependency in the medical image, and improving the global feature expression ability of the segmentation network. Our CiT-Net does not require a large number of labeled data for pre-training, but it can maximize the retention of local details and global semantic information in medical images. It has achieved the best segmentation performance on both dermoscopic images and liver datasets.

2 Method

2.1 Overall Architecture

The fusion of local and global features are clearly helpful for improving medical image segmentation. CNNs capture local features in medical images through convolution operation and hierarchical feature representation. In contrast, the Transformer network realizes the extraction of global features in medical images through the cascaded self-attention mechanism and the matrix operation with context interaction. In order to make full use of local details and global semantic features in medical images, we design a parallel interactive network architecture CiT-Net. The overall architecture of the network is shown in Figure 1 (a). CiT-Net fully considers the complementary properties of CNNs and Transformer. During the forward propagation process, CiT-Net continuously feeds the local details extracted by the CNNs to the decoder of the Transformer branch. Similarly, CiT-Net also feeds the global long-range relationship captured by the Transformer branch to the decoder of the CNNs branch. Obviously, the proposed CiT-Net provides better local and global feature representation than pure CNNs or Transformer networks, and it shows great potential in the field of medical image segmentation.

Specifically, CiT-Net consists of a patch embedding model, dynamically adaptive CNNs branch, cross-dimensional fusion Transformer branch, and feature fusion module. Among them, the dynamically adaptive CNNs branch and the crossdimensional fusion Transformer branch follow the design of U-Net and Swin-Unet, respectively. The dynamically adaptive CNNs branch consists of seven main stages. By using the weight coefficient and deformation offset adaptive DDConv in each stage, the segmentation network can better understand the local semantic features of medical images, better perceive the subtle changes of human organs or lesions, and improve the ability of extracting multi-scale change targets in medical images. Similarly, the cross-dimensional fusion Transformer branch also consists of seven main stages. By using (S)W-ACAM attention in each stage, as shown in Figure 1 (b), the segmentation network can better understand the global dependency of medical images to capture the position information between different organs, and improve the separability of the segmented object and the background in the medical images.

Although our CiT-Net can effectively improve the feature representation of medical images, it requires a large number of training data and network parameters due to the dualbranch structure. As the conventional Transformer network contains a lot of MLP layers, which not only aggravates the training burden of the network but also makes the number of model parameters rise sharply, resulting in the slow training of the model. Inspired by the idea of the Ghost network [Han et al., 2020], we redesign the MLP layer in the original Transformer and proposed a lightweight perceptron module (LPM). The LPM can help our CiT-Net not only achieve better medical image segmentation results than MLP but also greatly reduced the parameters and computational complexity of the original Transformer block, even the Transformer can achieve good results without a lot of labeled data training. It is worth mentioning that the dual-branch structure involves mutually symmetric encoders and decoders so that the parallel interaction network structure can maximize the preservation of local features and global features in medical images.

2.2 Dynamic Deformable Convolution

Vanilla convolution has spatial invariance and channel specificity, so it has a limited ability to change different visual modalities when dealing with different spatial locations. At the same time, due to the limitations of the receptive field, it is difficult for vanilla convolution to extract features of small targets or targets with blurred edges. Therefore, vanilla convolution inevitably has poor adaptability and weak generalization ability for complex medical images. Although the existing deformable convolution [Dai et al., 2017] and dynamic convolution [Chen et al., 2020] [Li et al., 2021] outperforms vanilla convolution to a certain extent, they still have the unsatisfied ability to balance the performance and size of networks when dealing with medical image segmentation.

In order to solve the shortcomings of current convolution operations, this paper proposes a new convolution strategy, DDConv, as shown in Figure 2. It can be seen that DDConv can adaptively learn the kernel deformation offset and weight coefficients according to the specific task and data distribution, so as to realize the change of both the shapes and the values of convolution kernels. It can effectively deal with the problems of large data distribution differences and large target deformation in medical image segmentation. Also, DDConv is plug-and-play and can be embedded in any network structure.

The shape change of the convolutional kernel in DDConv is based on the network learning of the deformation offsets. The segmentation network first samples the input feature map X using a square convolutional kernel S, and then performs a weighted sum with a weight matrix M. The square convolution kernel S determines the range of the receptive field, e.g., a 3 × 3 convolution kernel can be expressed as:

then the output feature map Y at the coordinate ϕn can be expressed as:

when the deformation offset △ϕm = {m = 1, 2, 3, . . . , N} is introduced in the weight matrix M, N is the total length of S. Thus the Equation (2) can be expressed as:

Through network learning, an offset matrix with the same size as the input feature map can be finally obtained, and the matrix dimension is twice that of the input feature map.

To show the convolution kernel of DDConv is dynamic, we first present the output feature map of vanilla convolution:

where σ is the activation function, W is the convolutional kernel weight matrix and y is the output feature map. In contrast, the output of the feature map of DDConv is:

where n is the number of weight coefficients, αn is the weight coefficients with learnable parameters and ˆ y is the output feature map generated by DDConv. DDConv achieves dynamic adjustment of the convolution kernel weights by linearly combining different weight matrices according to the corresponding weight coefficients before performing the convolution operation.

According to the above analysis, we can see that DDConv realizes the dynamic adjustment of the shape and weights of the convolution kernel by combining the convolution kernel deformation offset and the convolution kernel weight coefficient with a minimal number of calculation. Compared with directly increasing the number and size of convolution kernels, the DDConv is simpler and more efficient. The proposed DDConv not only solves the problem of poor adaptive feature extraction ability of fixed-size convolution kernels but also overcomes the defect that different inputs share the same convolution kernel parameters. Consequently, our DDConv can be used to improve the segmentation accuracy of small targets and large targets with blurred edges in medical images.

2.3 ShiftedWindowAdaptive Complementary Attention Module

The self-attention mechanism is the core computing unit in Transformer networks, which realizes the capture of longrange dependency of feature maps by utilizing matrix operations. However, the self-attention mechanism only considers the dependency in the spatial dimension but not the crossdimensional dependency between spatial and channels [Hong et al., 2021]. Therefore, when dealing with medical image segmentation with low contrast and high density noise, the self-attention mechanism is easy to confuse the segmentation targets with the background, resulting in poor segmentation results.

To solve the problems mentioned above, we propose a new cross-dimensional self-attention module called (S)WACAM. As shown in Figure 3, (S)W-ACAM has four parallel branches, the top two branches are the conventional dual attention module [Liu et al., 2021a] and the bottom two branches are cross-dimensional attention modules. Compared to popular self-attention modules such as spatial selfattention, channel self-attention, and dual self-attention, our proposed (S)W-ACAM can not only fully extract the longrange dependency of spatial and channels, but also capture the cross-dimensional long-range dependency between spatial and channels. These four branches complement each other, provide richer long-range dependency relationships, enhance the separability between the foreground and background, and thus improve the segmentation results for medical images.

The standard Transformer architecture [Dosovitskiy et al., 2020] uses the global self-attention method to calculate the relationship between one token and all other tokens. This calculation method is complex, especially in the face of highresolution and intensive prediction tasks like medical images where the computational costs will increase exponentially. In order to improve the calculation efficiency, we use the shifted window calculation method similar to that in Swin Transformer [Liu et al., 2021b], which only calculates the selfattention in the local window. However, in the face of our (S)W-ACAM four branches module, using the shifted window method to calculate self-attention does not reduce the overall computational complexity of the module. Therefore, we also designed the compact convolutional projection. First, we reduce the local size of the medical image through the shifted window operation, then we compress the channel dimension of feature maps through the compact convolutional projection, and finally calculate the self-attention. It is worth mentioning that this method can not only better capture the global high-dimensional information of medical images but also significantly reduce the computational costs of the module. Suppose an image contains h×w windows, each window size is M ×M, then the complexity of the (S)W-ACAM, the global MSA in the original Transformer, and the (S)W-MSA

in the Swin Transformer are compared as follows:

if the former term of each formula is a quadratic function of the number of patches hw, the latter term is linear when M is fixed (the default is 7). Then the computational costs of (S)W-ACAM are smaller compared with MSA and (S)WMSA.

Among the four parallel branches of (S)W-ACAM, two branches are used to capture channel correlation and spatial correlation, respectively, and the remaining two branches are used to capture the correlation between channel dimension C and space dimension H and vice versa (between channel dimension C and space dimension W). After adopting the shifted window partitioning method, as shown in Figure 2 (b), the calculation process of continuous Transformer blocks is as follows:

where ˆ T l and T l represent the output features of (S)WACAM and LPM, respectively. W-ACAM represents window adaptive complementary attention, SW-ACAM represents shifted window adaptive complementary attention, and LPM represents lightweight perceptron module. For the specific attention calculation process of each branch, we follow the same principle in Swin Transformer as follows:

where relative position bias B ∈ RM2×M2 , Q,K, V ∈ RM2× C8 are query, key, and value matrices respectively. C8 represents the dimension of query/key, andM 2 represents the number of patches.

After four parallel attention branches Out1, Out2, Out3 and Out4 are calculated, the final feature fusion output is:

where λ1, λ2, λ3 and λ4 are learnable parameters that enable adaptive control of the importance of each attention branch for spatial and channel information in a particular segmentation task through the back-propagation process of the segmentation network.

Different from other self-attention mechanisms, the (S)WACAM in this paper can fully capture the correlation between

† indicates the model is initialized with pre-trained weights on the ImageNet21K. "Para." refers to the number of parameters. "GFLOPs" is calculated under the input scale of 224 × 224. Since the dermoscopic images are 2D medical images, the comparison methods are all 2D networks.

spatial and channels, and reasonably use the context information of medical images to achieve long-range dependence modeling. Since our (S)W-ACAM effectively overcomes better feature representation of the defect that the conventional self-attention only focuses on the spatial self-attention of images and ignores the channel and cross-dimensional selfattention, it achieves the best image suffers from large noise, low contrast, and complex background.

2.4 Architecture Variants

We have built a CiT-Net-T as a base network with a model size of 11.58 M and a computing capacity of 4.53 GFLOPs. In addition, we built the CiT-Net-B network to make a fair comparison with the latest networks such as CvT [Wu et al., 2021] and PVT [Wang et al., 2021b]. The window size is set to 7, and the input image size is 224 × 224. Other network parameters are set as follows:

• CiT-Net-T: layer number = {2, 2, 6, 2, 6, 2, 2}, H = {3, 6, 12, 24, 12, 6, 3}, D = 96
• CiT-Net-B: layer number = {2, 2, 18, 2, 18, 2, 2}, H = {4, 8, 16, 32, 16, 8, 4}, D = 96,

D represents the number of image channels when entering the first layer of the dynamically adaptive CNNs branch and the cross-dimensional fusion Transformer branch, layer number represents the number of Transformer blocks used in each stage, and H represents the number of multiple heads in self-attention.

3 Experiment and Results

3.1 Datasets

We conducted experiments on the skin lesion segmentation dataset ISIC2018 from the International Symposium on Biomedical Imaging (ISBI) and the Liver Tumor Segmentation Challenge dataset (LiTS) from the Medical Image Computing and Computer Assisted Intervention Society (MICCAI). The ISIC2018 contains 2,594 dermoscopic images for training, but the ground truth images of the testing set have not been released, thus we performed a five-fold crossvalidation on the training set for a fair comparison. The LiTS contains 131 3D CT liver scans, where 100 scans of which are used for training, and the remaining 31 scans are used for testing. In addition, all images are empirically resized to 224 × 224 for efficiency.

3.2 Implementation Details

All the networks are implemented on NVIDIA GeForce RTX 3090 24GB and PyTorch 1.7. We utilize Adam with an initial learning rate of 0.001 to optimize the networks. The learning rate decreases in half when the loss on the validation set has not dropped by 10 epochs. We used mean squared error loss (MSE) and Dice loss as loss functions in our experiment.

3.3 Evaluation and Results

In this paper, we selected the mainstream medical image segmentation networks U-Net [Ronneberger et al., 2015], Attention Unet [Oktay et al., 2018], Swin-Unet [Cao et al., 2021], PVT [Wang et al., 2021b], CrossForm [Wang et al., 2021c] and the proposed CiT-Net to conduct a comprehensive comparison of the two different modalities datasets, ISIC2018 and the LiTS.

In the experiment of the ISIC2018 dataset, we made an overall evaluation of the mainstreammedical image segmentation network by using five indicators: Dice (DI), Jaccard (JA), Sensitivity (SE), Accuracy (AC), and Specificity (SP). Table 1 shows the quantitative analysis of the results of the proposed CiT-Net and the current mainstream CNNs and Transformer networks in the ISIC2018 dataset. From the experimental results, we can conclude that our CiT-Net has the minimum number of parameters and the lowest computational costs, and can obtain the best segmentation effect on the dermoscopic images without adding pre-training. Moreover, our CiT-Net-T network has only 11.58 M of parameters and 4.53 GFLOPs of computational costs, but still achieves the second-best segmentation effect. Our CiT-Net-B network,

† indicates the model initialized with pre-trained weights on ImageNet21K. "Para." refers to the number of parameters. "GFLOPs" is calculated under the input scale of 224 × 224. Compared with the comparison experiment oin the ISIC2018 dataset, 3D Unet and V-Net are introduced into the comparison experiment oin the LiTS-Liver dataset.

BAT, CvT, and CrossForm have similar parameters or computational costs, but in the ISIC2018 dataset, the division Dice value of our CiT-Net-B is 1.02%, 3.00%, and 3.79% higher than that of the BAT, CvT, and CrossForm network respectively. In terms of other evaluation indicators, our CiT-Net-B is also significantly better than other comparison methods.

In the experiment of the LiTS-Liver dataset, we conducted an overall evaluation of the mainstream medical image segmentation network by using five indicators: DI, VOE, RVD, ASD and RMSD. Table 2 shows the quantitative analysis of the results of the proposed CiT-Net and the current mainstream networks in the LiTS-Liver dataset. It can be seen from the experimental results that our CiT-Net has great advantages in medical image segmentation, which further verifies the integrity of CiT-Net in preserving local and global features in medical images. It is worth noting that the CiTNet-B and CiT-Net-T networks have achieved good results in medical image segmentation in the first and second place, with the least number of model parameters and computational costs. The division Dice value of our CiT-Net-B network without pre-training is 1.20%, 1.03%, and 1.01% higher than that of the Swin-Unet, TransUNet, and CvT network with pre-training. In terms of other evaluation indicators, our CiTNet-B is also significantly better than other comparison methods.

3.4 Ablation Study

In order to fully prove the effectiveness of different modules in our CiT-Net, we conducted a series of ablation experiments on the ISIC2018 dataset. As shown in Table 3, we can see that the Dynamic Deformable Convolution (DDConv) and (Shifted) Window Adaptive Complementary Attention Module ((S)W-ACAM) proposed in this paper show good performance, and the combination of these two modules, CiT-Net shows the best medical image segmentation effect. At the same time, the Lightweight Perceptron Module (LPM) can significantly reduce the overall parameters of the CiT-Net.

4 Conclusion

In this study, we have proposed a new architecture CiTNet that combines dynamically adaptive CNNs and crossdimensional fusion Transformer in parallel for medical image segmentation. The proposed CiT-Net integrates the advantages of both CNNs and Transformer, and retains the local details and global semantic features of medical images to the maximum extent through local relationship modeling and long-range dependency modeling. The proposed DDConv overcomes the problems of fixed receptive field and parameter sharing in vanilla convolution, enhances the ability to express local features, and realizes adaptive extraction of spatial features. The proposed (S)W-ACAM self-attention mechanism can fully capture the cross-dimensional correlation between feature spatial and channels, and adaptively learn the important information between spatial and channels through network training. In addition, by using the LPM to replace the MLP in the traditional Transformer, our CiT-Net significantly reduces the number of parameters, gets rid of the dependence of the network on pre-training, avoids the challenge of the lack of labeled medical image data and easy over-fitting of the network. Compared with popular CNNs and Transformer medical image segmentation networks, our CiT-Net shows significant advantages in terms of operational efficiency and segmentation effect.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grants 62271296, 62201334 and 62201452, in part by the Natural Science Basic Research Program of Shaanxi under Grant 2021JC-47, and in part by the Key Research and Development Program of Shaanxi under Grants 2022GY-436 and 2021ZDLGY08-07.