[Transformer]TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation

Medical Image Segmentation)

Abstract

医学图像分割是众多临床诊断的先决任务，目前主要基于CNN框架，然而CNN的局限性在于无法捕获长程依赖关系，目前的一些优化办法在于通过搭建深层次编码网络提取深层特征，但这样又会导致网络冗余、丢失本地 细节信息。
本文则提出一种并行多分支结构-TransFuse，将Transformer与CNN并行的结合在一起，这样可以同时捕获全局依赖关系了局部空间信息。
此外还提出一种新的融合方式-BiFusion可以有效的从多分支中融合多尺度特征。在2D和3D图像中均验证了TransFuse的有效性，同时网络参数显著降低、推理速度明显提升。

Section I Introduction

CNN在诸多医学图像分割任务中取得了优异性能，如多器官分割、肝部病变分割、脑补分割等，显示出CNN在建模特定任务特征表示的强大能力；但CNN的一大问题是在捕获全文上下文信息方面缺乏效率，如果通过堆叠、扩大感受野的方式则需要连续的下采样-卷积操作，使得网络结构十分深，这一过程也会导致局部信息的丢失，而拒不细节信息对于密集预测任务也十分重要。


在小型数据集上 较深的网络层次还会导致训练不稳定、过拟合等问题。有的研究探索了使用non-local 注意力机制来建模全局信息，但是往往计算复杂度是输入的平方项，并不适合高分辨率的图像。
Transformer最初用于NLP中的seq2seq任务，ViT则是首个纯基于self-attention的Transformer模型用于处理图像问题，但是需要在大型数据集ImageNet预训练之后的结果；SETR则是将编解码网络中的encoder部分替换成了Transforemr，在自然图像分割任务中达到了SOTA。但Transformer也有自身局限，就是无法较好的建模细粒度特征，尤其是对于医学图像，细节特征十分重要。因此SETR这种纯Transformer模型用于医学图像分割效果并不理想，因为其在建模局部信息时缺少空间的归纳偏执。

为了将CNN与Transformer物尽其用，目前正致力于将CNN与Transformer结合，比兔TransUNet则是先试用CNN提取低层次细节特征，然后送入Transformer提取全局交互信息；借助skip connection进行信息交互在多器官CT图像分割任务中创造了新的记录；但是前人的工作大多聚焦于将卷积模块替换为Transformer或者将二者堆叠在一起，为了进一步发挥二者的作用，本文提出一种新的组织架构-TransFuse-会并行的进行CNN特征编码和Transformer分割网络。然后在经过本文提出的BiFusion模块进行特征融合。


TransFuse有以下优势：
（1）可以有效的捕获低层次空间特征和高级语义特征；

（2）无需搭建深层次网络，减轻梯度消失、特征无法有效复用等问题；


（3）大幅度提升了模型效率和推理速度，提升了部署在云端或终端的效率。


TransFuse是首个并行的将CNN与Transformer结合在一起的网络架构，实验结果显示在息肉分割任务中达到了SOTA。




Section II Proposed Method

Fig 1展示了TransFuse结构示意图，包含两个并行分支：
（1）CNN分支
会逐渐增加感受野尺寸，提取不同尺度特征

（2）Transformer分支
首先计算全局注意力最后回复细节信息
两个分支同层同分辨率的特征会送入BiFusion模块进行特征融合，主要通过SA和双线性哈达玛内积来有选择性的进行信息融合。

通过基于门控的skip connection融合不同尺度的feature map获得最终的分割图。


TransFuse这样搭建有两大优势：


一是可以有效的发挥CNN和Transformer各自的优势，在不用搭建深层次网络的同时有效捕获全局信息，还能保持精确的低层细节信息；


二是BiFusion模块可以在融合过程中同时利用CNN和Transformer提取的有效特征，使得融合的特征更加紧凑、表达能力更强。
Transformer Branch.

Transformer分支的设计遵循其编码-解码结构，对于HXWX3的输入首先会切成patch（patch number= 16）;随后patch经过线性映射并展平，在送入Transformer之前还会嵌入可训练的位置信息。
Transformer中包含L层MSA和MLP，SA的计算表示为：




经过Transformer encoder处理后的结果会送到Decoder，在解码器部分则采用渐进上采样的方式（PUP），类似SETR的操作。
首先会将Z^L的输出reshape到原始二维维度，拥有D0个通道；然后使用两个连续的上采样-卷积层来恢复空间分辨率，最终获得不同分辨率的上采样结果，用于与CNN的feature map进行特征融合。

CNN branch.

在CNN网络中通过逐渐进行下采样操作获得不同尺度的特征图，这样十分耗费资源，也会使得网络层次很深。为了更好的发挥Transformer的特长，本文将CNN流水线的下采样操作替换为Transformer，使用Transformer来提取全局上下文信息。
这样一方面可以减少网络深度一方面也可以保留更丰富的低层细节信息。
比如搭建5层特征提取模块，可以保留前四层的提取结果，可以只保留前4层的输出，然后将他们与Transformer的结果融合获得融合后的特征提取。

本文的CNN分支设计十分灵活，可以用于任何现成的CNN网络。


BiFusion Module



BiFusion模块用于融合来自CNN和Transformer编码后的特征。包含的操作有：
通道注意力、空间注意力、3x3卷积、残差连接。


哈达玛内积主要为了建模两个分支的细粒度关联；经过残差连接后有效融合了当前分辨率的全局和局部的上下文信息。为了获得最终的分割图还使用了AG skip-connection。
Loss Function
损失函数包括了加权mIoU损失和BCE损失，

在边界处的像素赋予了更大的权重。



最终的分割图是由一个simply head生成，直接将输入的feature map调整会原始分辨率，通过卷积生成M张分割图，M为类别数目。



最终的损失函数表述为：





Section III Experiments

Datasets


Polyp Segmentation


息肉分割
使用了4个息肉分割数据集：Kvasir [14], CVC-ClinicDB , CVC-ColonDB ,
EndoScene,ETIS



training:testing = 1450:798




皮肤损伤数据集




ISIC2017





training:validation:testing = 2000:150:600





髋关节分割
Hip segmentation





前列腺分割







Results

为了验证TransFuse的有效性，本文在2D,3D图像上均开展了实验。
息肉分割实验中没有使用数据增强；皮肤和关节分割则使用了随机旋转、水平翻转等数据增强


TransFuse也设置了3种规格


TransFuse-S：对标的是ResNet34,Deit-S


TransFuse-L:对标的是ResNet50,Dei-B


TransFuse-L*:对标的是Resnetv2-50和ViT-B



Polyp Segmentation




Table1展示了息肉分割的对比结果，可以看到TransFuse-S/L与其他CNN网络相比均达到了SOTA，而且参数量比PraNet等减少了20%,实时性也更好;而预训练过的TransFuse-L*的性能也优于SETR和TransUNet。




Skin Lesion Segmentation





在皮肤损伤分割实验中性能指标是Jaccard指数和Dice系数以及逐像素点的accuracy。从Table 2可以看出，TransFuse性能优于UNet++,同时不需要任何预处理或后处理，而UNet++需要使用预训练的R34作为主干。





Hip Segmentation






Table 3展示的是髋关节分割的对比结果，主要需要分割腓骨、左股骨、右股骨。与UNet++和HRNet相比，TransFuse在HD和ASD这两个指标参数上均更加优秀，充分证明了本文提出的TransFuse可以有效捕获精细的结构，生成更加清晰准确的轮廓。






Prostate Segmentation






nnUNet是目前在前列腺分割排名第一的分割网络，Table 4则是TransFuse与nnUMet的对比结果。可以看到与nnUNet-3d相比，TransFuse-S不仅性能更好，而且参数量减少了41%,吞吐率增加了50%。







Ablation Study

Table 5展示的是并行分支的消融实验，Table 6展示的是BiFusion的消融实验。
可以看到两个分支选择CNN搭配Transformer性能是最佳的，BiFusion结合空间注意力、通道注意力、计算内积都会提升性能。

Section IV Conclusion

本文提出一种将CNN与Transformer结合的新的组合策略-晚融合-用于医学图像分割。TransFuse结构通过CNN分支可以有效建模空间关系，Transformer分支可以有效建立长程关系，在诸多分割任务中均达到了SOTA，并且计算更加高效，参数量更少，推理速度更快。

本文希望可以为基于Transformer的结构提供一种新的视角。未来本文计划将进一步提升原始transformer层的效率。同时测试TransFuse在其他医学相关任务上的性能，如病灶检测、疾病分类等。