[论文解读] TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation

Yundong Zhang, Huiye Liu+1 more|arXiv (Cornell University)|Feb 16, 2021

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

一句话总结

TransFuse 介绍了一种并行的 CNN 和 Transformer 架构，以及 BiFusion 融合模块，能够联合捕捉低层细节与全局上下文用于医学图像分割，在参数更少、推理更快的条件下实现了最先进的结果。

ABSTRACT

Medical image segmentation - the prerequisite of numerous clinical needs - has been significantly prospered by recent advances in convolutional neural networks (CNNs). However, it exhibits general limitations on modeling explicit long-range relation, and existing cures, resorting to building deep encoders along with aggressive downsampling operations, leads to redundant deepened networks and loss of localized details. Hence, the segmentation task awaits a better solution to improve the efficiency of modeling global contexts while maintaining a strong grasp of low-level details. In this paper, we propose a novel parallel-in-branch architecture, TransFuse, to address this challenge. TransFuse combines Transformers and CNNs in a parallel style, where both global dependency and low-level spatial details can be efficiently captured in a much shallower manner. Besides, a novel fusion technique - BiFusion module is created to efficiently fuse the multi-level features from both branches. Extensive experiments demonstrate that TransFuse achieves the newest state-of-the-art results on both 2D and 3D medical image sets including polyp, skin lesion, hip, and prostate segmentation, with significant parameter decrease and inference speed improvement.

研究动机与目标

• 解决 CNN 在建模全局上下文和保留局部细节方面在医学图像分割中的局限性。
• 提出一种并行的 CNN-Transformer 架构（TransFuse），配备 BiFusion 模块用于高效特征融合。
• 证明使用 BiFusion 的并行分支在参数更少、推理更快的条件下仍能获得强大的分割性能。
• 在多样的 2D/3D 医学数据集（息肉、皮肤病变、髋部、前列腺）上展示最先进的结果。

提出的方法

• 两条并行分支以不同方式处理信息：一个 CNN 分支强调局部空间细节，另一个 Transformer 分支建模全局上下文。
• BiFusion 模块使用通道注意力和空间注意力以及 Hadamard 乘积融合来自两个分支的多层特征。
• 采用基于注意力的跳跃连接进行晚期融合，聚合融合后的特征以生成分割图。
• 损失函数结合加权 IoU 与加权二元交叉熵，并在各分支之间进行深度监督。
• 变体（TransFuse-S、TransFuse-L、TransFuse-L*）探索不同的骨干网络选择（CNN + DeiT/ViT 骨干）。

实验结果

研究问题

• RQ1 并行的 CNN 与 Transformer 分支是否能够在不需要非常深的网络的情况下捕捉用于医学图像分割的互补信息？
• RQ2 BiFusion 模块在融合多等级 CNN 与 Transformer 特征方面的效果如何？
• RQ3 提出的 TransFuse 变体是否在多样的 2D/3D 医学分割任务中提高了准确性和效率？
• RQ4 TransFuse 在息肉、皮肤病变、髋部和前列腺分割数据集上的性能如何？

主要发现

• TransFuse 在多项 2D 和 3D 医学分割任务（息肉、皮肤病变、髋部、前列腺）上取得了最先进的结果。
• 通过利用并行的 Transformer 分支，减少了单分支 CNN 的下采样，从而得到整体更浅的模型并提升全局上下文。
• TransFuse-S 在大约 26.3M 参数下达到有竞争力或更优的结果，推理速度仍然很高（例如在 RTX 2080 Ti 上达到 98.7 FPS）。
• TransFuse-L* 在具有挑战性的息肉数据集上进一步优于 TransUnet，同时保持良好的效率（例如大约 45.3 FPS）。
• ISIC 2017 皮肤病变结果显示 TransFuse-S 取得了强劲的指标（例如带预训练骨干的 Jaccard/Dice 大约为 0.795/0.872）。
• 消融研究表明 BiFusion 模块和并行分支设计相对于顺序融合与简单拼接策略对性能提升有贡献。