The Fully Convolutional Transformer for Medical Image Segmentation

论文：The Fully Convolutional Transformer for Medical Image Segmentation (thecvf.com)

代码：Thanos-DB/FullyConvolutionalTransformer (github.com)

期刊/会议：WACV 2023

摘要

我们提出了一种新的transformer，能够分割不同形态的医学图像。医学图像分析的细粒度特性所带来的挑战意味着transformer对其分析的适应仍处于初级阶段。UNet压倒性的成功在于它能够欣赏分割任务的细粒度性质，这是现有的基于transformer的模型目前不具备的能力。为了解决这个缺点，我们提出了全卷积transformer(FCT)，它建立在卷积神经网络学习有效图像表示的能力的基础上，并将它们与transformer的能力相结合，有效地捕获其输入中的长期依赖关系。FCT是医学影像文献中第一个全卷积Transformer模型。它分两个阶段处理输入，首先，它学习从输入图像中提取长期语义依赖关系，然后学习从特征中捕获分层的全局属性。FCT结构紧凑、准确、健壮。我们的结果表明，它在不需要任何预训练的情况下，在不同数据形态的多个医学图像分割数据集上，大大优于所有现有的transformer架构。FCT在ACDC数据集上比其直接模型高出1.3%，在Synapse数据集上高出4.4%，在Spleen数据集上高出1.2%，在ISIC 2017数据集上高出1.1%，在dice metric上的参数少了多达5倍。在ACDC Post-2017MICCAI-Challenge在线测试集上，我们的模型在未见过的MRI测试用例上设置了新的最先进的技术，优于大型集成模型以及参数更少的nnUNet。

1、简介

医学图像分割是计算机辅助诊断的关键工具。它有助于检测和定位图像中病变的边界，有助于快速识别肿瘤和癌变区域的潜在存在。这有可能加快诊断，提高检测肿瘤的可能性，并使临床医生更有效地利用他们的时间，对患者的结果有利[15]。传统上，现代医学图像分割算法构建为对称的自顶向下编码器-解码器结构，首先将输入图像压缩(编码)到潜在空间，然后学习解码图像中感兴趣区域的位置。将中间信号的水平传播(跳越连接)添加到这个垂直信息流中，我们就得到了UNet架构，这可以说是最近分割算法中最具影响力的飞跃。今天大多数现代分割系统都包括UNet或其变体。UNet成功的关键在于其全卷积的性质。UNet在其结构中不估计任何非卷积可训练参数。

基于卷积神经网络(CNN)的UNet模型在医学图像分割任务中的准确性和性能方面取得了巨大的成功。然而，为了真正帮助临床医生进行早期疾病诊断，它们仍然需要额外的改进。卷积算子固有的局部特性是CNN的一个关键问题，因为它阻止了它们利用来自输入图像的长范围语义依赖。人们提出了各种方法来为CNN添加全局上下文，最引人注目的是引入注意力机制，以及扩大卷积核以增加核的感受野。然而，这些方法都有自己的缺点。Transformer在语言学习任务中取得了巨大的成功，因为它们能够有效地处理非常长范围的序列依赖。这导致它们最近适应了各种视觉任务[7,18,21,22]。最近提出的架构，如ViT[7]，已经超过了cnn在基准成像任务上的性能，而最近对ViT的许多改进，如CvT [36]， CCT[10]和Swin Transformer[25]，已经表明transformer不需要庞大的数据消耗模型，甚至可以处理少量数据，从而超过CNN的性能。通常，ViT风格的模型首先从图像中提取离散的非重叠patch(在NLP中称为token)。然后，他们通过位置编码将空间定位注入到这些patch中，并将此表示通过标准transformer层来建模数据中的长期语义依赖关系。

考虑到CNN和Transformer的明显优点，我们认为医学图像分割的下一步是一个完全卷积编码器-解码器深度学习模型，能够有效地利用医学图像中的长期语义依赖。为了实现这一目标，我们提出了第一个用于医学图像分割的全卷积Transformer。我们新颖的全卷积Transformer层构成了我们模型的主要构建块。它包含两个关键组件，一个卷积注意力模块和一个全卷积Wide-Focus模块(见第3节)。我们将我们的贡献形式化如下:

1. 我们提出了第一个用于医学图像分割的全卷积Transformer，它超越了所有现有的基于卷积和transformer的医学图像分割架构的性能，用于多个二分类和语义分割数据集。
2. 我们提出了一种新型的全卷积transformer层，它使用卷积注意力模块来学习长范围的语义上下文，然后通过宽焦点模块使用多分辨率空洞卷积创建分层的局部到全局上下文。
3. 通过广泛的消融研究，我们展示了我们模型的各种构建块在其对模型性能影响的背景下的影响。

2、相关工作

早期的CNN和Attention模型：UNet[29]是第一个用于医学图像分割的CNN模型。最早将注意力模型引入医学图像分割的工作之一，是通过将门控函数应用于UNet[26]的编码器到解码器的特征传播。FocusNet[17]等方法采用双编码器-解码器结构，其中注意力门选学习将相关特征从一个UNet的解码器传播到下一个UNet的编码器。FocesNet++[19]是在分组卷积的各种过滤器组中集成注意力机制的第一个作品之一。还有许多UNet的变体，它们使用不同的残差块来增强特征提取[32,28,33,20,16]。UNet++[43]在编码器和解码器之间创建了嵌套的分层密集跳过连接路径，以减少它们之间学习特征的语义差距。作为最近最具影响力的UNet变体，nnUNet[14]自动调整自身来预处理数据，并选择最适合任务的最佳网络架构，而不需要人工干预。

Transformer模型:最初的Transformer架构[31]彻底改变了自然语言处理任务，并迅速成为视觉理解任务的模型[7]。Transformer在视觉方面工作得很好，因为它们能够创建长范围的视觉环境，但存在固有的缺点，不能利用CNN等图像中的空间环境。最近的工作转向了克服这一缺陷的可能解决方案。CvT [36]， CCT[10]和Swin Transformer[25]都是在transformer中集成足够的空间环境的尝试。在医学图像分割中，大多数现有研究着眼于创建用于特征处理的Transformer-CNN混合模型。与Attention UNet[26]类似，UNet Transformer[27]增强了CNN，在跳过连接内增加了多头注意。TransUNet[5]是最早提出的用于医学图像分割的Transformer-CNN混合模型之一，它使用Transformer编码器馈送到级联卷积解码器。与TransUNet类似，UNETR[12]和Swin UNETR[11]在编码器上使用Transformer和卷积解码器来构造分割地图。Transfuse[40]运行双分支编码器，一个带有卷积层，另一个带有transformer层，并将其特征与新颖的BiFusion模块结合起来。然而，这个模型的解码器是卷积的。

当前的工作：最近有一个转变，从创建混合Transformer-CNN模型，到改进transformer块本身，以处理医学图像的细微差别。Swin UNet[3]是第一个提出用于处理医学图像的纯transformer的架构。这里的纯指的是仅由transformer层提取和处理的图像特征，而不需要预训练的骨干网络架构。DS-TransUNet[24]引入Transformer Interactive Fusion模块，以获得更好的表示全局依赖。这两个模型的计算核心都是Swin Transformer块。同时进行的工作，如nnFormer[42]和DFormer[37]，试图利用医学图像中的本地和全局上下文，通过特别制作的多头自我关注块来满足这一任务。这些模型的主要缺点是它们固有的注意力投射和特征处理的线性性质，FCT旨在缓解这一点。

现有的医学影像分割模型目前至少存在以下三个局限性之一。它们要么基于CNN主干网络，要么使用卷积层创建，因此限制了它们超越感受野以获得图像语义上下文的能力(早期CNN方法)。他们试图将Transformer集成到他们的特征处理管道中，以利用它们创建长期语义上下文的能力，但反过来，使模型庞大且计算复杂(混合Transformer-CNN)。他们试图通过创建用于分割的纯Transformer模型来减少这种计算负担，而不试图在低级特征提取阶段(并发工作)对局部空间上下文建模。与现有方法不同，我们的全卷积Transformer没有这些缺点，同时仍然是一个纯基于Transformer的医疗图像分割架构。补充资料中的表4额外总结了FCT与现有模型相比的主要差异。

3、全卷积Transformer

给定一个数据集 { X , Y } \{\mathbf{X}, \mathbf{Y}\} {X,Y}，其中，$\mathbf{X}$是我们模型的输入图像，$\mathbf{Y}$是相应的语义或二分类分割映射。对于每个图像${\mathbf{x}}_i\in {\mathbb{R}}^{H\times W\times C}$，其中$H$和$W$为图像的空间分辨率， C = { 3 , … , N } C =\{3,\ldots, N\} C={3,…,N}为输入通道数，我们的模型产生一个输出分割映射${\mathbf{y}}_i\in {\mathbb{R}}^{H\times W\times K}$，其中， K ∈ { 1 , … , D } K∈\{1,\ldots,D\} K∈{1,…,D}。FCT的输入是从输入3D图像的每个切片中采样的2D patch。我们的模型遵循熟悉的UNet形状，FCT层作为其基本构建块。与现有的方法不同，我们的模型既不是CNN-Transformer的混合，也不是Transformer-UNet的结构，它使用现成的transformer层来编码或细化输入特征。它首先从图像中提取重叠的patch，然后创建基于patch的扫描嵌入，然后在这些patch上应用多头自注意，从而构建特征表示。然后通过我们的Wide-Focus模块处理给定图像的输出投影，以从投影中提取细粒度信息。图1显示了我们的网络体系结构的概述。

3.1 FCT层

每个FCT层都从LayerNormalization-Conv-Conv-Maxpool操作开始。我们从经验上注意到，与直接先创建图像的patch-wise投影相比，在3×3内核大小较小的patch上按顺序应用这些连续卷积有助于更好地编码图像信息。每个卷积层后面都有一个Gelu激活函数。我们的FCT块与其他模型块不同的第一个实例是通过它对医学成像的卷积注意力应用。

MaxPool的输出被输入到转换函数$\mathbf{T}(\cdot )$中，转换函数$\mathbf{T}(\cdot )$将其转换为新的token映射。我们选择的$\mathbf{T}(\cdot )$是Depthwise-Convolution operator。我们选择一个较小的内核大小3×3, 步长为s×s和一个有效的填充，以确保:(1)与大多数现有工作不同，提取的patch是重叠的，并且(2)卷积操作不会始终改变输出大小。接下来是LayerNormalization操作。得到的token映射$p_{i+1}\in {\mathbb{R}}^{W_t\times H_t\times C_t}$被平化为$W_t{H}_t\times C_t$，创建了我们的patch嵌入式输入。下一个例子是，我们的FCT层不同于现有的基于transformer的医学成像应用方法，是通过它的注意力投影。所有现有模型都采用线性逐点线性映射来进行多头自我注意(MHSA)计算。这导致Transformer模型失去空间信息，这对成像应用非常重要。现有的方法试图通过卷积增强来缓解这个问题，使其适应成像任务。然而，这为所提出的模型增加了额外的计算成本。受[36]中提出的方法的启发，我们将MHSA块中的逐点线性映射替换为Depthwise-Convolution，以降低计算成本，并从图像中利用更好的空间上下文信息。patch嵌入和卷积注意力投影构成了我们的卷积注意力的组成部分。与[36]不同的是，我们注意到用LayerNormalization替换BatchNormalization有助于提高性能。此外，删除Point-wise Convolution会导致一个更简单的模型，而不会损失任何性能。Depthwise-Convolution提供的空间上下文进一步消除了对位置编码的需求，位置编码用于在输入中插入空间信息，并顺序跟踪每个patch的位置，从而进一步简化了架构设计。

一般的Transformer层遵循线性层MHSA块，因此丢失了图像中的所有空间上下文。直接用卷积替换这些线性层是一种相对简单的方法，可以缓解这个问题并提高性能。然而，医学图像需要细粒度的信息处理。记住这一点，我们采用了一个多分支卷积层，其中一层对MHSA输出应用空间卷积，而其他层应用空洞卷积，增加感受野，以获得更好的空间上下文。然后，我们通过求和来融合这些特征，并将它们传递到特征聚合层。这种特征聚合是通过另一个空间卷积算子完成的。我们称这个模块为Wide-Focus。残差连接用于增强整个层的特征传播。最后的特征被重新塑造，并进一步传播到下一个FCT层。图1(上)显示了FCT层。

3.2 编码器

我们的模型的编码器包含四个FCT层，负责特征提取和传播。对于第$l$个transformer层，卷积注意力模块的输出为，${\mathbf{z}}_l^{\prime }=\mathbf{MHSA}(z_{l-1})+{\mathbf{z}}_{l-1}^{q/k/v}$，其中，${\mathbf{z}}_{l-1}^{q/k/v}=\mathbf{Flatten}(\mathbf{DepthConv}(\mathbf{Reshape}(z_{l-1})))$。$\mathbf{MHSA}({\mathbf{z}}_{l-1})=softmax(Q{K}^T/\sqrt{d})V$。然后由Wide-Focus (WF)模块处理$z_l^{\prime }$, ${\mathbf{z}}_l=\mathbf{WF}({\mathbf{z}}_l)+{\mathbf{z}}_l^{\prime }$。我们进一步为编码器注入金字塔风格的图像输入，目的是在不同尺度上突出显示不同类别和更小的ROI特征。值得注意的是，即使没有这种多尺度图像金字塔输入，我们的模型也能够获得最先进的结果。数据的(瓶颈)潜在编码是使用另一个FCT层创建的。

3.3 解码器

解码器将瓶颈表示作为其输入，并学习从该信息中重新采样二分类或语义分割映射。为了在解码器层中创建更好的上下文相关性，还使用从编码器到解码器的跳过连接，其中来自编码器层的具有相同分辨率的特征映射与解码器层连接。解码器的形状与编码器对称。解码器中的层对应于编码器中的图像金字塔层，输出中间分割映射，提供额外的监督并提高模型的预测能力。上下文相关性是通过首先对特征量进行上采样，然后将其传递到FCT层以了解其最佳可能表示来创建的。我们没有在FCT的最低规模上采用深度监管，因此我们的模型不是“完全深度监管”。这是因为我们观察到，输入图像扫描中的感兴趣区域(roi)有时太小，无法在最低尺度(28 × 28)进行分割，这导致模型性能较差。这种低规模的输出在模型中增加了强烈的偏差，以预测一些输出roi作为背景类。

4、实验

数据集：(MRI) Automatic Cardiac Diagnosis Challenge (ACDC) [2], (CT) Synapse Multi-organ Segmentation Challenge1, (CT) Spleen Segmentation Dataset [1] and (Dermoscopy) ISIC 2017 [6] Skin Cancer Segmentation Challenge。

实验细节：模型输入有两种尺寸224 x 224，384 x 384。Adam优化器，学习率1e-3。

5、结果

6、总结

我们提出了全卷积transformer，它能够准确地执行二分类和语义分割任务，参数比现有模型更少。FCT在参数数量上比nnFormer小5倍以上，比TransUNet和LeViT-UNet小3倍以上。FCT层由两个关键组件组成——卷积注意力和Wide-Focus。卷积注意力通过使用深度可分离卷积为模型创建重叠的patch，消除了在patch创建阶段对位置编码的需求。我们基于深度可分离卷积的MHSA块集成了空间信息，首次在医学成像背景下估计长距离语义依赖关系。从我们的消融实验中可以看出，Wide-Focus有助于利用医学图像中存在的细粒度特征信息，并且是提高transformer块性能的重要因素。我们通过在多个高度竞争的不同模式和维度的细分数据集上的最先进的结果证明了我们模型的能力。我们的FCT块是第一个为医学成像应用而提出的全卷积transformer块，并且可以轻松扩展到医学成像的其他领域和应用。我们相信我们的模型可以作为未来分割任务的有效骨干网络，并为基于transformer的医学图像处理的创新铺平道路。