[论文解读] TEC-Net: Vision Transformer Embrace Convolutional Neural Networks for Medical Image Segmentation

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

Advanced Neural Network Applications被引用 5

一句话总结

TEC-Net 提出了一种混合医学图像分割网络，其在卷积神经网络（CNN）分支中引入了动态可变形卷积（DDConv），以实现自适应的局部特征提取；在 Transformer 分支中引入了一种新型（移位）窗口自适应互补注意力模块（(S)W-ACAM），以建模跨维度的长距离依赖关系。该方法在无需预训练的情况下，以更少的参数量和更低的推理成本实现了最先进（SOTA）的性能表现。

ABSTRACT

The hybrid architecture of convolution neural networks (CNN) and Transformer has been the most popular method for medical image segmentation. However, the existing networks based on the hybrid architecture suffer from two problems. First, although the CNN branch can capture image local features by using convolution operation, the vanilla convolution is unable to achieve adaptive extraction of image features. Second, although the Transformer branch can model the global information of images, the conventional self-attention only focuses on the spatial self-attention of images and ignores the channel and cross-dimensional self-attention leading to low segmentation accuracy for medical images with complex backgrounds. To solve these problems, we propose vision Transformer embrace convolutional neural networks for medical image segmentation (TEC-Net). Our network has two advantages. First, dynamic deformable convolution (DDConv) is designed in the CNN branch, which not only overcomes the difficulty of adaptive feature extraction using fixed-size convolution kernels, but also solves the defect that different inputs share the same convolution kernel parameters, effectively improving the feature expression ability of CNN branch. Second, in the Transformer branch, a (shifted)-window adaptive complementary attention module ((S)W-ACAM) and compact convolutional projection are designed to enable the network to fully learn the cross-dimensional long-range dependency of medical images with few parameters and calculations. Experimental results show that the proposed TEC-Net provides better medical image segmentation results than SOTA methods including CNN and Transformer networks. In addition, our TEC-Net requires fewer parameters and computational costs and does not rely on pre-training. The code is publicly available at https://github.com/SR0920/TEC-Net.

研究动机与目标

• 为解决标准卷积网络在医学图像分割中固定感受野和参数共享带来的局限性。
• 克服传统 Transformer 中自注意力机制在医学图像应用中对通道和跨维度注意力的忽视。
• 在保持小样本医学数据集高分割精度的前提下，降低模型复杂度并消除对预训练的依赖。
• 有效整合 CNN 和 Transformer 分支，实现局部与全局特征学习的互补。

提出的方法

• 在 CNN 分支中引入动态可变形卷积（DDConv），通过学习空间可变的卷积核权重和偏移量，实现任务自适应的特征提取。
• 在 Transformer 分支中提出 (S)W-ACAM 模块，联合建模空间、通道和跨维度注意力，提升长距离依赖关系的建模能力。
• 采用紧凑的卷积投影层，降低 Transformer 编码器中的计算成本和参数数量。
• 在 Transformer 前馈网络中，用可学习参数模块（LPM）替代标准 MLP，以减小模型规模并消除对预训练的依赖。
• 设计了带有跳跃连接的双分支 U 形编码器-解码器架构，融合 CNN 和 Transformer 分支的特征，以提升定位和分割性能。
• 通过深层网络层实现多尺度特征融合与渐进式优化，提升轮廓和目标边界的分割精度。

实验结果

研究问题

• RQ1 与标准卷积相比，动态可变形卷积是否能提升医学图像分割中的局部特征表示能力？
• RQ2 在 Transformer 中联合建模跨维度（空间-通道）注意力，是否能提升对低对比度和含噪复杂医学图像的分割性能？
• RQ3 混合 CNN-Transformer 架构是否能在参数更少、计算成本更低的前提下实现 SOTA 性能？
• RQ4 将标准 MLP 替换为可学习参数模块（LPM）在多大程度上可消除医学图像分割中对预训练的需求？
• RQ5 CNN 与 Transformer 分支的融合在多大程度上增强了局部细节保持与全局上下文建模能力？

主要发现

• 在 ISIC2018 数据集上，TEC-Net 的 Dice 分数达到 90.72%，优于包括 Swin-Unet 和 U-Net 变体在内的现有 SOTA 方法。
• 消融实验表明，结合 DDConv、(S)W-ACAM 和 LPM 可获得最佳性能，当所有组件均启用时，Dice 分数达 90.88%。
• TEC-Net 将模型参数量减少至 11.58M，显著低于 U-Net+Swin-Unet 基线的 46.92M，展现出极高的参数效率。
• 网络在无需预训练的情况下仍保持高精度，表现为仅使用 LPM 时 Dice 分数为 88.43%，而结合 (S)W-ACAM 和 LPM 时达 89.46%，表明其在小样本数据集上的鲁棒性。
• 可视化结果表明，网络通过解码器中特征融合逐步优化分割结果，尤其在边界和轮廓细节区域表现突出。
• (S)W-ACAM 模块能有效捕捉跨维度依赖关系，表现为注意力图更聚焦于病灶边界和器官结构。