IJCAI 2023|CiT-Net: Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation
论文标题: CiT-Net: Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation
论文发表会议: IJCAI 2023
论文地址: https://arxiv.org/abs/2306.03373
论文代码: https://github.com/SR0920/CiT-Net
关键词: 动态可变形卷积;滑窗自适应互补注意力;卷积神经网络;Transformer
一、摘要
卷积神经网络(CNN)和Transformer的混合架构已经成为医学图像分割的最流行方法。然而,现有的混合架构存在两个问题。首先,尽管CNN分支能够利用卷积运算捕获图像的局部特征,但是常规卷积无法实现对图像特征的自适应提取。其次,尽管Transformer分支能够对图像进行全局建模,但是常规的自注意力仅关注图像的空间自注意力而忽略了通道与跨维度自注意力,因此对具有复杂背景的医学图像分割精度较低。为了解决这些问题,我们提出了Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation (CiT-Net)。我们的网络有两个优点。首先,在CNN分支中设计了动态可变形卷积(Dynamic Deformable Convolution, DDConv),不仅解决了固定尺寸卷积核自适应特征提取能力差的问题,而且解决了不同输入共享相同卷积核参数的缺陷,有效提升了对医学图像的特征表达能力。其次,在Transformer分支中设计了滑窗自适应互补注意力(Shifted Window Adaptive Complementary Attention Module, SW-ACAM)与紧致卷积投影(compact convolutional projection),使网络在使用极少参数量和计算量的情况下充分学习医学图像的跨维度长程依赖关系。实验表明,与包括CNN和Transformer网络在内的SOTA方法相比,所提出的CiT-Net提供了更好的医学图像分割结果。此外,CiT-Net保持着更少的参数量与计算量,并且不依赖于预训练。
二、引言
图像分割是指将图像分割成若干个特定的、具有独特性质的区域。在广泛的应用场景中起着至关重要的作用。医学图像分割是这一领域的一个重要应用,其为临床治疗提供了很多好处,显著提高了临床治疗的诊断效率与准确性。同时医学图像分割算法提供的可视化结果不仅能够提供对人体异常区域的检测,而且可以用于指导临床医生。因此,准确的医学图像分割已经成为了计算机辅助诊疗、患者病情分析、图像引导手术、组织器官重建和制定治疗规划的关键组成部分。相比于普通的RGB图像,医学图像通常存在着噪声大、对比度低和边缘模糊等问题,所以如何快速、准确的从医学图像中分割出特定的人体器官和病灶,并提取出图像的关键特征,一直是一项巨大的挑战。
早期出现的传统医学图像分割算法是基于医学专家使用专业知识设计的手工特征,这类方法具有较强的数学基础与理论支撑,但是这些算法对于人体不同器官或者病灶的泛化性较差。后来受到全卷积网络(fully convolutional networks, FCN)和编解码器的启发,设计出了最早应用于医学图像分割的U-Net[3]网络。该网络提出之后,其对称的U型编解码结构得到了人们的广泛关注。同时由于U-Net模型小、分割效果好,使得深度学习在医学图像分割中取得了突破性的进展。随之启发了一系列基于U-Net结构改进的医学图像分割网络,例如2D的U-Net++,ResDO-UNet,2.5D的RIU-Net,3D的3D Unet,V-Net等。卷积神经网络(CNN)在医学图像分割领域的迅速发展,很大程度上归功于卷积运算。因为卷积运算能够通过分层特征提取的方式捕获与表征医学图像信息。然而由于卷积运算中感受野的局限性,尽管CNN在提取图像的局部特征方面有着巨大的优势,但是在捕捉图像的全局特征方面CNN却存在着明显的缺陷。
针对CNN在获取医学图像全局特征方面存在的缺陷,学者们提出了能够捕获医学图像全局特征的Transformer网络架构。Transformer通过复杂的空间变换和长距离关系建模,构成了对图像信息的全局表征,有效解决了CNN仅能获取图像局部特征的缺陷。将Transformer应用于医学图像分割,代表性方法如Swin-Unet,BAT,Swin UNETR,UCTransNet。这些方法大致可以分为纯Transformer架构和CNNs与Transformer的混合架构。纯Transformer的架构基于自注意力实现了对医学图像的长程依赖建模。但是由于Transformer缺乏归纳偏置(Inductive bias),导致Transformer在类似于医学图像这样的小规模数据集中不能得到广泛的应用。同时因为Transformer容易忽略局部细节特征,从而降低了医学图像中一些微小病灶、尺度变化大的目标与背景之间的可分性。而CNNs与Transformer的混合架构,通过利用CNNs与Transformer优势互补的特点,实现了对医学图像的局部与全局信息建模,从而取得了更好的医学图像分割效果。然而这些网络仍然存在以下两个问题:首先,这些网络在对图像的局部特征建模时忽略了器官的形变和病灶的不规则问题,导致对形变较大的器官和病灶的局部特征表达能力较弱。其次,这些网络在对图像进行全局特征建模时忽略了特征图空间与通道之间的相关性,导致自注意力表达不充分。针对上述问题,我们主要贡献如下:
- 提出了一种新的动态可变形卷积 (DDConv),DDConv通过任务自适应学习,可以灵活改变卷积自身的权重系数和形变偏置。DDConv能够克服标准卷积及其变体卷积Atrous convolution, Involution等感受野固定、卷积核参数共享的问题,同时增强了对医学图像局部特征的表达能力,实现了对空间特征的自适应提取,提高了对医学图像中微小病灶与尺度变化大的目标的感知能力。
- 提出了一种新的滑窗自适应互补注意力(SW-ACAM),SW-ACAM通过权重系数自适应学习的四个平行分支,实现了对医学图像跨维度的全局建模。相比于目前流行的注意力机制,例如:CBAM,Non-Local。SW-ACAM充分弥补了常规注意力机制中对空间与通道之间跨维度关系建模不足的缺陷,其能够捕获医学图像中跨维度的长距离关联特征,增强了对医学图像中分割目标与背景之间的可分性。相比于CNN和原始Transformer网络显著降低了网络的参数量和计算量,同时能够实现对医学图像全局特征的捕获。
- 提出了一种基于动态自适应CNN与跨维度特征融合 Transformer的并行新网络结构用于医学图像分割,称为CiT-Net。相比于目前流行的CNNs与Transformer混合架构,例如Swin-Unet,Swin UNETR。CiT-Net通过利用CNNs和Transformer并行交互的方式紧密结合不同分辨率下的局部与全局特征,来增强表征学习(representation learning),最大限度的保留医学图像中的局部特征和全局特征。值得注意的是,CiT-Net不但舍弃了预训练,而且拥有着更少的参数量和计算量,分别是11.58M和4.53GFLOPs。
三、方法
3.1 网络结构
融合图像的局部特征和全局特征是医学图像准确分割的重要内容。CNN通过卷积运算以分层特征提取的方式捕获医学图像中的局部特征。相比之下,Transformer网络通过级联的自注意力机制,利用具有上下文交互的矩阵运算,实现对医学图像中全局特征的提取。为了充分利用医学图像中的局部细节特征与全局语义特征,我们设计了一种并行交互的网络架构CiT-Net,网络的整体架构如图1(a)所示。CiT-Net网络的设计充分考虑了CNN和Transformer两种不同风格网络的互补性。在网络前向传播过程中,我们不断地将CNN分支提取到的局部细节信息反馈至Transformer分支的解码部分,用来丰富Transformer分支的局部特征。同样,我们也将Transformer分支捕获到的全局长程建模关系反馈至CNN分支的解码部分,用来增强CNN分支对于全局特征的感知能力。显然,与单独的CNN或者Transformer网络相比较,所提出的CiT-Net充分继承了CNN和Transformer的结构优势与泛化优势,更好的提供了医学图像数据中的局部与全局特征表示,展现出了在医学图像分割领域的巨大潜力。
具体来说,CiT-Net主要由patch embedding,动态自适应 CNN分支,跨维度融合Transformer分支以及特征融合四部分组成。其中动态自适应 CNN分支和跨维度融合Transformer分支分别遵循了U-Net和Swin-Unet的设计思路。动态自适应 CNN分支主要由7个阶段组成,通过在每阶段使用权重系数和形变偏置自适应的DDConv卷积,分割网络能够更好的理解医学图像的局部语义特征,更好的感知人体器官或病灶微妙的变化,提高了对医学图像中多尺度变化目标的提取能力。类似地,跨维度融合Transformer分支也主要由7个阶段组成,通过在每个阶段使用(S)W-ACAM注意力,如图1(b)所示,分割网络能够更好的理解医学图像的全局依赖关系用于捕获人体不同器官之间的位置信息,提高了对医学图像中分割目标与背景的可分性。常规的Transformer网络之所以需要大量的数据进行训练,是因为其包含了大量的MLP层。大量使用MLP层不仅加剧了网络的训练负担,而且使得模型参数量急剧上升,最终导致模型训练缓慢,需要大量的标注数据进行预训练才能取得好的效果。我们受到Ghost轻量化的思想,将原始Transformer中的MLP层进行了重新设计,提出了lightweight perceptron module(LPM),通过LPM在Transformer块的大量使用,不但取得了比MLP更好的医学图像分割效果,而且大幅度减小了原始Transformer块的参数量与复杂度,使得Transformer不需要大量的标注数据训练也能取得好的成绩。值得一提的是,双分支结构是由相互对称的编解码器组成,这样并行交互的网络结构,能够最大限度的保留医学图像中的局部特征与全局特征。
3.2 动态可变形卷积
常规卷积由于空间不变性(spatial invariance)和通道特异性(channel specificity),剥夺了卷积核应对不同空间位置时对不同视觉模式的变化能力。同时由于感受野的局限性,导致常规卷积对小目标或者边缘模糊的目标特征提取困难。因此在面对边缘模糊、目标形变大的医学影像分割任务时,常规卷积对于未知的变化适应性差,泛化能力弱。而现有的可变形卷积和动态卷积虽然在一定程度上弥补了常规卷积的一些缺陷,但在应对医学影像分割问题时,其仍然存在网络特征表达能力不足与参数量巨大难以平衡的问题。
为解决目前卷积存在的缺陷,本文提出了动态可变形卷积。如图2所示,DDConv可以根据特定的医学图像分割任务和数据分布,自适应地学习卷积核形变偏移量和卷积核权重系数,同时通过网络的反向传播实现端到端的训练,从而实现卷积核形状与权重的双重变化,因此可以有效应对医学图像分割任务中不同目标之间数据分布差异大、目标形变大的问题。同时,DDConv是即插即用的,可以嵌入于任何网络结构中。
根据上述分析,我们可以看到DDConv通过以最小的计算量将卷积核形变偏移量和卷积核权重系数进行组合,实现了卷积核形状和权重的动态调整。与直接增加卷积核的数量和大小相比,DDConv更简单、更高效。所提出的DDConv不仅解决了固定大小卷积核自适应特征提取能力差的问题,而且解决了不同输入共享相同卷积核参数的缺陷,因此可以用于提高医学图像中小目标和边缘模糊的大目标的分割精度。
3.3滑窗自适应互补注意力
自注意力机制是Transformer中的核心计算单元,其通过利用具有上下文交互的矩阵运算实现了对长距离相关特征的捕获。但是自注意力机制在运算过程中,只考虑了空间维度中的依赖关系,未考虑空间与通道之间跨维度的相互依赖关系,缺乏跨维度间的长程依赖建模。因此在解决具有比度低、噪声大等复杂背景的医学图像分割任务时,自注意力机制容易将分割目标与图像背景混淆,导致目标分割困难,器官或病变部位边缘分割结果粗糙。
为了解决上述问题,我们提出了一种新的跨维度自注意力模块称为(S)W-ACAM。如图3所示,(S)W-ACAM有四个平行分支,顶部两个分支是传统的双注意力模块,底部两个分支为跨维度注意力模块。与流行的自注意力模块(如空间自注意力、通道自注意力和双重自注意力)相比,我们提出的(S)W-ACAM不仅可以充分提取空间和通道的长程依赖关系,还可以捕捉空间和通道之间的跨维度长程依赖关系。这四个分支相互补充,提供了更丰富的长距离依赖关系,增强了前景和背景之间的可分性,从而提高了医学图像的分割效果。
不同于其他的自注意力机制,本文的(S)W-ACAM能够充分捕获通道与空间之间的相关性,合理利用医学图像的上下文信息进行长程依赖建模,弥补了常规自注意力仅关注图像的空间自注意力而忽略了通道与跨维度自注意力的缺陷,对医学影像中对比度低、噪声大等背景复杂的分割目标有着更加精确的特征提取能力,从而提升了分割网络对医学图像的整体特征表达能力。
四、实验
在ISIC2018数据集的实验中,我们通过采用Dice (DI),Jaccard (JA),Sensitivity (SE),Accuracy (AC),Specificity (SP)五项指标对主流的医学图像分割网络进行了总体评估。表1展示了本文提出的CiT-Net与目前的主流的CNNs和Transformer网络在ISIC2018数据集中的结果定量分析。从实验结果可以得出,我们提出的CiT-Net拥有最少的模型参数量与计算量,同时在不需要添加预训练的情况下,能够获得最佳的皮肤镜图像分割效果。其中,CiT-Net-T网络仅有11.58M的参数量和4.53GFLOPs的计算量,却仍然取得第二名的分割效果。CiT-Net-B网络与BAT、CvT、CrossForm有着类似的参数量或者计算量,但在ISIC2018数据集中,CiT-Net-B的分割Dice数值比BAT、CvT、CrossForm网络分别高1.02%、3.00%、3.79%。在其他评估指标方面,CiT-Net-B也明显优于其他对比方法。
在LiTS-Liver数据集的实验中,我们通过采用DI,VOE,RVD,ASD,RMSD五项指标对主流的医学影像分割网络进行了总体评估。表2展示了本文提出的CiT-Net与目前的主流网络在LiTS-Liver数据集中的结果定量分析。通过实验结果可以看出,本文提出的CiT-Net在医学图像分割方面的巨大优势,进一步验证了CiT-Net对医学图像中的局部特征和全局特征保留的完整性。值得注意的是,本文提出的CiT-Net-B和CiT-Net-T网络,在拥有最少的模型参数量与计算量的前提下,仍然获得了第一名和第二名的医学图像分割好成绩。不带预训练的CiT-Net-B网络的分割Dice数值比带有预训练的Swin-Unet、TransUNet和CvT网络分别高1.20%、1.03%和1.01%。在其他评估指标方面,CiT-Net-B也明显优于其他对比方法。
五、消融实验
为了充分证明本文提出的CiT-Net中不同模块的有效性,我们在ISIC2018数据集上进行了一系列的消融实验。如表3所示,可以看出本文提出的DDConv和SW-ACAM在医学图像分割方面都能展现出很好的性能,将两者结合之后的CiT-Net更是表现出了最佳的医学图像分割效果。
六、结论
在本文中,我们提出了一种动态自适应CNN和跨维度融合 Transformer并行结合的新架构CiT-Net用于医学图像分割领域,CiT-Net整合了CNNs和Transformer中各自独特的优势,通过局部关系建模和长程依赖建模,对医学图像中的局部细节特征和全局语义特征得到了最大限度的保留。提出的DDConv卷积策略克服了常规卷积感受野固定、参数共享的问题,增强了对医学影像局部特征的表达能力,实现了对空间特征的自适应提取。提出的SW-ACAM自注意力机制能够充分捕获医学图像通道与空间之间跨维度的相关性,通过网络训练自适应的学习空间与通道之间的重要信息。此外,我们通过引入LPM轻量化策略代替了传统Transformer中的MLP,显著降低了传统CNNs和Transformer网络中的参数量,摆脱了网络对预训练的依赖,解决了医学图像中标注数据少、网络容易过拟合的问题。相比于目前主流的CNNs和Transformer医学图像分割网络,本文的CiT-Net在运行效率和分割效果方面都展现出了显著的优势。