[医学图像Transformer]TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

论文地址

开源代码

摘要

传统U-Net通常在显式建模长期依赖（long-range dependency）方面表现出局限性。Transformers，设计用于序列到序列的预测，已经作为具有固有的全局自注意机制的替代架构出现，但由于低层次细节不足，可能导致有限的定位能力。在本文中，我们提出了具有Transformers和U-Net优点的TransUNet，作为医学图像分割的一个强有力的替代方案。一方面，Transformer将来自卷积神经网络(CNN)特征图的标记图像编码为输入序列，用于提取全局上下文。另一方面，解码器对编码后的特征进行上采样，然后将其与高分辨率CNN特征图相结合，实现精确定位。

存在的问题及解决方案

由于卷积运算固有的局域性，基于CNN的方法通常在建模显式的长期关系时表现出局限性。因此，这些结构通常表现出较弱的性能，特别是对于在纹理、形状和大小方面表现出患者间较大差异的目标结构。与之前的基于CNN的方法不同，Transformer不仅在建模全局上下文方面很强大，而且在大规模的预训练下，对下游任务也表现出优越的可移植性。

在这篇论文中，我们提出了第一探索在医学图像分割的背景下Transformer的潜力的研究。然而，有趣的是，我们发现单纯的使用(即使用一个Transformer对标记图像进行编码，然后直接将隐藏的特征表示上采样为全分辨率的密集输出)不能产生令人满意的结果。

这是由于Transformer将输入视为一维序列，在所有阶段只专注于建模全局上下文，因此导致低分辨率特征，缺乏详细的定位信息。而直接上采样到全分辨率无法有效地恢复这些信息，从而导致粗分割结果。另一方面，CNN的架构(如U-Net)提供了一种提取低级视觉线索的方法，可以很好地弥补这些空间细节。

为此，我们提出医学图像分割框架TransUNet，该框架从序列到序列预测的角度建立了自注意机制。为了弥补Transformer带来的特征分辨率的损失，TransUNet采用了一种混合CNN- transformer架构，以利用CNN特征和Transformer编码的全局上下文的详细高分辨率空间信息。受U-Net结构的启发，Transformer编码的自注意特征被上采样，并与编码路径跳过的不同高分辨率CNN特征相结合，以实现精确定位。我实证结果表明，与之前基于CNN的自注意方法相比，基于Transformer的架构提供了更好的方法来利用自注意。此外，我们观察到，更密集的结合低级特征通常会导致更好的分割精度。

方法

本文的方法基于ViT（An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale），这是一篇非常强大的Transformer在CV上的应用，有兴趣的可以看看，核心在于patch的使用，本文也使用了这张patch的方法，即将图像分块再输入到Transformer中。

这篇文章的结构其实简单，了解Transformer和U-Net一看就懂，这里主要的一个创新在于，不是直接利用Transformer当成编码器，而是在Transformer前面还用了一组CNN做前置编码器，这样做的好处在于可以提取CNN中的高分辨率特征图加入到解码过程中（毕竟Transformer是展成了一维序列，丢失了定位信息），来弥补单纯Transformer的定位问题。

实验结果

注：CUP表示类似U-Net中skip-connection的级联结构

从这张表中可以发现，单纯使用ViT做解码器效果并不如ResNet50，原因在于Transformer能够很好地捕获有利于分类任务的高级语义，但缺乏用于分割医学图像精细形状的低级线索。 而将Transformer与CNN结合的TransUnet则能够有效缓解这个问题，从而实现新的性能。

从图中可以发现单纯的基于CNN的U-Net和AttnUNet方法更有可能对器官进行过分割(如第二行，脾被AttnUNet过分割而被UNet欠分割)，这表明基于Transformer的模型，如我们的TransUNet或R50-ViT-CUP有更强大的能力来编码全局上下文和区分语义。第一行的结果显示，我们的TransUNet预测的误报比其他方法少，这表明TransUNet在抑制这些噪声预测方面比其他方法更有优势。在基于Transformer的模型中进行比较，我们可以观察到，在边界和形状方面，R50-ViT-CUP的预测往往比TransUNet的更粗糙（例如第二行中关于胰腺的预测）。此外，在第三行，TransUNet正确预测左右肾，而R50-ViT-CUP错误填充左肾内孔。这些观察结果表明，TransUNet能够进行更精细的分割，并保留详细的形状信息。原因是TransUNet既享有高级全局上下文信息的好处，也享有低级细节的好处，而R50-ViT-CUP只依赖高级语义特性。

结论

在本文中，为了充分利用Transformer的力量，TransUNet被提出，它不仅通过将图像特征作为序列来编码强全局上下文，而且通过U型混合结构设计很好地利用了低级CNN特征。作为医学图像分割的替代框架，TransUNet取得了优于各种竞争方法的性能，包括基于CNN的自注意方法。