[论文解读] CiT-Net: Convolutional Neural Networks Hand in Hand with Vision Transformers for Medical Image Segmentation

Tao Lei, Rui Sun+4 more|arXiv (Cornell University)|Jun 6, 2023

Advanced Neural Network Applications参考文献 38被引用 5

一句话总结

该论文提出CiT-Net，一种用于医学图像分割的混合CNN-Transformer架构，其在CNN分支中引入了动态可变形卷积，在Transformer分支中采用了移位窗口自适应互补注意力模块。该模型在无需预训练的情况下，以更少的参数量和FLOPs实现了最先进（SOTA）的性能，优于Swin-Unet和TransUNet等模型，在ISIC2018和LiTS-Liver数据集上表现更优。

ABSTRACT

The hybrid architecture of convolutional neural networks (CNNs) and Transformer are very popular for medical image segmentation. However, it suffers from two challenges. First, although a CNNs branch can capture the local image features using vanilla convolution, it cannot achieve adaptive feature learning. Second, although a Transformer branch can capture the global features, it ignores the channel and cross-dimensional self-attention, resulting in a low segmentation accuracy on complex-content images. To address these challenges, we propose a novel hybrid architecture of convolutional neural networks hand in hand with vision Transformers (CiT-Net) for medical image segmentation. Our network has two advantages. First, we design a dynamic deformable convolution and apply it to the CNNs branch, which overcomes the weak feature extraction ability due to fixed-size convolution kernels and the stiff design of sharing kernel parameters among different inputs. Second, we design a shifted-window adaptive complementary attention module and a compact convolutional projection. We apply them to the Transformer branch to learn the cross-dimensional long-term dependency for medical images. Experimental results show that our CiT-Net provides better medical image segmentation results than popular SOTA methods. Besides, our CiT-Net requires lower parameters and less computational costs and does not rely on pre-training. The code is publicly available at https://github.com/SR0920/CiT-Net.

研究动机与目标

• 解决标准CNN在医学图像分割中固定感受野和参数共享带来的局限性。
• 克服标准Transformer在通道和跨维度注意力方面的不足，从而提升在复杂医学图像上的性能。
• 通过用轻量级感知机模块替代MLP，降低模型复杂度并消除对预训练的依赖。
• 在参数量和计算成本低于现有混合CNN-Transformer模型的前提下，实现更高的分割精度。
• 开发一种鲁棒、高效且可泛化的医学图像分割架构，以同时保留局部与全局特征。

提出的方法

• 在CNN分支中引入动态可变形卷积（DDConv），通过学习卷积核偏移量实现自适应的空间特征提取，避免固定卷积核的限制。
• 在Transformer分支中提出移位窗口自适应互补注意力模块（(S)W-ACAM），通过自适应注意力机制建模空间特征与通道特征之间的跨维度依赖关系。
• 采用轻量级感知机模块（LPM）替代Transformer中的标准MLP，显著降低模型参数量和计算成本。
• 设计双分支架构，使CNN与Transformer分支并行处理特征，有效融合局部与全局表征。
• 使用紧凑的卷积投影层，在保持效率的同时支持Transformer分支中的长距离依赖建模。
• 端到端训练模型，无需预训练，降低在有限医学数据集上的过拟合风险，提升泛化能力。

实验结果

研究问题

• RQ1 与标准CNN相比，动态可变形卷积是否能提升医学图像分割中的局部特征表示？
• RQ2 在Transformer分支中引入跨维度注意力是否能提升复杂医学图像的分割精度？
• RQ3 轻量级感知机模块是否能有效替代Transformer中的标准MLP，从而在不损失性能的前提下降低模型大小和计算成本？
• RQ4 在无需预训练的情况下，所提出的混合架构在分割精度、参数量和FLOPs方面与SOTA方法相比如何？
• RQ5 自适应注意力与动态卷积的融合在多大程度上提升了医学图像中局部与全局特征的保留能力？

主要发现

• CiT-Net-B在ISIC2018数据集上达到90.88%的Dice分数，优于Swin-Unet、TransUNet和CvT分别1.20%、1.03%和1.01%，且无需预训练。
• CiT-Net-T在ISIC2018上取得90.72%的最佳Dice分数，仅需11.58M参数和4.53 GFLOPs，是目前最高效的SOTA模型。
• 在LiTS-Liver数据集上，CiT-Net-B在全部五项指标（DI、VOE、RVD、ASD、RMSD）上表现最优，展现出强大的泛化能力和鲁棒性。
• 消融实验表明，结合DDConv与(S)W-ACAM可获得最佳性能，两者同时使用时Dice分数相比基线提升1.2%。
• LPM将最佳配置下的模型参数量降低至9.67M，显著减少计算成本，同时保持高精度。
• 在ISIC2018上，CiT-Net的Dice分数分别优于BAT、CvT和CrossForm 1.02%、3.00%和3.79%，且参数量更少、FLOPs更低。