[论文解读] ConvFormer: Plug-and-Play CNN-Style Transformers for Improving Medical Image Segmentation

Xian Lin, Zengqiang Yan+3 more|arXiv (Cornell University)|Sep 9, 2023

Advanced Neural Network ApplicationsComputer Science被引用 3

一句话总结

本文提出 ConvFormer，一种即插即用的 CNN 风格变换模块，通过用池化、CNN 风格自注意力（CSA）和卷积前馈网络替代标准视觉变换器组件，缓解医学图像分割中的注意力坍塌问题。ConvFormer 直接在 2D 特征图上操作，支持自适应、核可扩展的注意力机制，提升了长距离依赖学习能力，并在 ACDC、ISIC 和 ICH 数据集上，对多种最先进基于变换器的模型实现了稳定性能提升。

ABSTRACT

Transformers have been extensively studied in medical image segmentation to build pairwise long-range dependence. Yet, relatively limited well-annotated medical image data makes transformers struggle to extract diverse global features, resulting in attention collapse where attention maps become similar or even identical. Comparatively, convolutional neural networks (CNNs) have better convergence properties on small-scale training data but suffer from limited receptive fields. Existing works are dedicated to exploring the combinations of CNN and transformers while ignoring attention collapse, leaving the potential of transformers under-explored. In this paper, we propose to build CNN-style Transformers (ConvFormer) to promote better attention convergence and thus better segmentation performance. Specifically, ConvFormer consists of pooling, CNN-style self-attention (CSA), and convolutional feed-forward network (CFFN) corresponding to tokenization, self-attention, and feed-forward network in vanilla vision transformers. In contrast to positional embedding and tokenization, ConvFormer adopts 2D convolution and max-pooling for both position information preservation and feature size reduction. In this way, CSA takes 2D feature maps as inputs and establishes long-range dependency by constructing self-attention matrices as convolution kernels with adaptive sizes. Following CSA, 2D convolution is utilized for feature refinement through CFFN. Experimental results on multiple datasets demonstrate the effectiveness of ConvFormer working as a plug-and-play module for consistent performance improvement of transformer-based frameworks. Code is available at https://github.com/xianlin7/ConvFormer.

研究动机与目标

• 解决将视觉变换器应用于小规模医学图像数据集时出现的注意力坍塌问题。
• 改善基于变换器的医学图像分割模型中自注意力图的收敛性和多样性。
• 设计一种即插即用模块，无需架构重构即可增强纯变换器和 CNN-Transformer 混合架构。
• 通过将自注意力建模为可学习的卷积核，实现可扩展的自适应长距离特征学习。
• 在多种数据集和主干网络架构上展示一致的性能提升。

提出的方法

• 用 2D 卷积和最大池化替代 ViT 的 1D tokenization，以保留空间局部性并降低分辨率。
• 提出 CNN 风格自注意力（CSA），将自注意力矩阵构建为自适应、可扩展的卷积核，实现像素级长距离依赖建模。
• 采用卷积前馈网络（CFFN）在 CSA 后细化特征，保持 2D 空间结构。
• 使用可学习的缩放因子 α 控制 CSA 中的感受野大小，实现局部与全局上下文之间的动态适应。
• 直接在 2D 输入图像上处理，避免标准 ViT 中使用的 1D 序列转换。
• 设计为即插即用模块，可轻松集成到现有基于变换器的分割框架中。

Figure 1: Visualization of attention maps from the selected layers of the first head in different transformer frameworks. The darker the color, the closer the dependency.

实验结果

研究问题

• RQ1 CNN 风格的变换器设计能否缓解低数据量医学图像分割中的注意力坍塌？
• RQ2 与 1D 序列相比，在 2D 特征图上操作是否能提升注意力多样性并改善模型收敛性？
• RQ3 基于自适应、核可扩展注意力的即插即用模块能否在多种基于变换器的分割架构中持续提升性能？
• RQ4 CSA 中感受野的可扩展性如何影响医学图像基准上的分割精度？
• RQ5 所提出的模块是否能在不同数据集和主干模型上实现医学图像分割中的泛化？

主要发现

• 当集成到纯变换器主干网络 SETR 中时，ConvFormer 在 ACDC 数据集上实现了平均 Dice 分数 3.86% 的提升。
• 在 ISIC 数据集上，使用 SETR 时 ConvFormer 将 Dice 分数提升了 1.38%，使用 Patcher 时提升了 1.07%。
• 在 ICH 数据集上，使用 SETR 时 ConvFormer 将 Dice 分数提升了 1.39%，使用 Patcher 时提升了 1.15%，表明在多个数据集上均实现了稳定增益。
• 在 TransUNet 和 TransFuse 等 CNN-Transformer 混合模型中，尽管 CNN 在训练中占主导地位，ConvFormer 仍提供了稳定的性能提升。
• 消融实验确认 α = 0.4 时性能最佳（ACDC 上 Dice 达 91.00%），表明感受野缩放达到最优。
• 可视化结果表明，与基线模型相比，ConvFormer 生成的注意力图更具多样性与结构性，有效缓解了注意力坍塌。

Figure 2: Comparison between vanilla vision transformer and ConvFormer. CBR is short for the combination of convolution, batch normalization, and Relu. Multiple heads are omitted for simplicity.