赞
踩
Transformer作为自然语言处理的首选模型,在医学影像界引起了很少的关注。考虑到利用长期依赖关系的能力,Transformer有望帮助非典型卷积神经网络克服其空间归纳偏差的固有缺点。然而,大多数最近提出的基于Transformer的分割方法只是将Transformer作为辅助模块来帮助将全局上下文编码为卷积表示。为了解决这个问题,我们引入了nnFormer(not-another Transformer),这是一种用于3D医学图像分割的3D Transformer。nnFormer不仅利用交织卷积和自注意操作的结合,还引入了局部和全局基于体积的自注意机制来学习体积表示。此外,nnFormer提出在U-Net类架构中使用跳跃注意来代替传统的跳跃连接中的连接/求和操作。实验表明,nnFormer在三个公共数据集上显著优于先前基于Transformer的同类工作。与nnUNet相比,nnFormer产生的HD95明显较低和DSC结果相当。此外,我们还证明了nnFormer和nnUNet在模型集成方面是高度互补的。
Transformer已经成为自然语言处理(NLP)问题的现实选择,最近同样在基于视觉的应用中被广泛利用,其背后的核心思想是应用自注意力机制来捕获长距离的依赖关系。相比CNN网络(convnets),Transformer解决了局部性的归纳偏差,使其更有能力处理非局部的相互作用。也有调查显示,与convnets相比,transformers的预测误差与人类的预测误差更加一致。
鉴于Transformer比convnets具有天然的优势,有许多方法试图将Transformer应用于医学图像分析领域。
Chen等人[11]首次提出了TransUNet,以探索Transformer在医学图像分割方面的潜力。TransUNet的整体架构与U-Net[12]相似,其中卷积网络作为特征提取器,而Transformer帮助编码全局背景。事实上,TransUNet及其大多数类似版本(TransFuse、MeT、TransClaw
U-Net、TransAttUnet)的一个主要特点是将convnets作为主体,在此基础上进一步应用Transformer来捕捉长距离的依赖关系。然而,这样的特点可能会导致一个问题,那就是Transformer的优势没有得到充分的发挥。换句话说,我们认为一到两层的Transformer不足以将长期依赖关系与卷积表征纠缠在一起,而卷积表征往往包含精确的空间信息,并提供层次化的概念。
为了解决上述问题,一些研究者开始使用Transformer作为分割模型的主干。【Convolution-free medical image segmentation using transformers】首次引入了一个无卷积的分割模型,将扁平化的图像表示转发给Transformer,然后将其输出重组为与分割Mask相一致的三维张量。最近,Swin Transformer表明,通过参考卷积网络中使用的特征金字塔,Transformer可以通过对特征图进行适当的下采样来学习不同尺度上的分层概念。受此启发,SwinUNet用分层的Transformer块在一个类似U-Net的架构中构建编码器和解码器,在此基础上,DS-TransUNet又增加了一个编码器,以接受不同大小的输入。SwinUNet和DS-TransUNet都实现了对TransUNet的持续改进。尽管如此,他们并没有探索如何适当地结合卷积和自注意力来构建一个最佳的医学分割网络。
相比之下,nnFormer(即not-another transFormer)使用的是混合主干网络,其中卷积和自注意力被交错使用,以充分发挥其优势。图1展示了nn-Former的编码器中使用的不同组件的效果。
首先,我们把一个轻量级的卷积嵌入层放在Transformer块的前面。与[17:无卷积]中直接平整原始像素和应用一维预处理相比,卷积嵌入层编码精确的(即像素级)空间信息,并提供低水平但高分辨率的三维特征。在嵌入块之后,Transformer和卷积下采样块交织在一起,将长期依赖关系与不同尺度的高层次和分层物体概念完全混合在一起,这有助于提高学习表征的泛化能力和鲁棒性。
nnFormer的另一个贡献在于提出了一种具有计算效率的方法来利用切片间依赖关系。具体来说,nnFormer提出联合使用基于局部体积的多头自我注意(LV-MSA)和基于全局体积的多头自我注意(GV-MSA)来构建特征金字塔,并为在局部和全局三维体积上学习表征提供足够的感受野,然后将其汇总以进行预测。
与原始的多头自我注意(MSA)相比,所提出的策略可以大大降低计算的复杂性,同时产生有竞争力的分割性能。此外,受机器翻译任务中使用的注意力机制的启发,我们引入了跳过注意力来取代类似U-Net架构的跳过连接中的非典型串联/求和操作,这进一步提高了分割结果。
总而言之,我们的贡献可以概括为以下几点。
在这一节中,我们主要回顾了借助Transformer来改善医学图像分割结果的方法。由于他们中的大多数采用了卷积和自注意力的混合结构,我们根据大多数干系是卷积的还是基于Transformer的,将它们分为两类。
TransUNet首次应用Transformer来改善医学图像的分割结果。TransUNet将卷积网络作为一个特征提取器,为输入的slice切片生成一个特征图。然后,patch Embedding被应用于瓶颈中的特征图的patch,而不是ViT中的原始图像。同时,与TransUNet类似,Li等人[Medical
image segmentation using squeeze-and-expansion transformers]提出使用挤压注意力块来规范Transformer的自注意力模块,并使用扩展块来学习眼底图像的多样化表示,这些都是在convnets内的瓶颈中实现的。TransFuse引入了BiFusion模块,将基于convnet的浅层编码器和基于transformer的分割网络的特征进行融合,对二维图像进行最终预测。与TransUNet相比,TransFuse主要将自注意力机制应用于输入嵌入层,以改善二维图像的分割模型。
Yun等人采用Transformer纳入光谱信息,与卷积特征编码的光谱信息融合在一起,以解决高光谱病理的问题。Xu等人广泛研究了Transformer和卷积网之间的权衡,并提出了一个更有效的编码器,名为LeViT-UNet。Li等人[ More than encoder: Introducing transformer decoder to upsample]提出了一种新的上采样方法,并将其纳入UNet的解码器,以模拟长期依赖关系和全局信息,从而获得更好的重建结果。TransClaw U-Net在UNet中利用了具有更多卷积特征金字塔的Transformer。TransAttUNet[16]探索了应用Transformer自我注意与卷积全局空间注意的可行性。CoTr采用Transformer来捕捉来自不同层卷积网络的多尺度卷积特征的长期依赖性。TransBTS首先利用三维卷积网络来提取体积空间特征,并对输入的三维图像进行下采样以产生分层表示。然后,TransBTS中编码器的输出被重塑为矢量(即tokens),并送入Transformer进行全局特征建模,之后附加一个普通的卷积解码器来上采样特征图以达到重建的目的。
与这些直接采用卷积网作为特征提取器的方法不同,我们的nnFormer在功能上依赖于卷积和基于Transformer的模块,它们交错使用,以发挥彼此的优势。
Valanarasu等人提出了一个门控的轴向注意力模型(即MedT),该模型通过在自注意力中引入一个总的控制机制,扩展了现有的卷积网络的架构。Karimi等人取消了卷积运算,建立了一个基于Transformer的三维分割模型。其主要思想是首先将局部体积块分割成三维patch块,然后将其扁平化并嵌入到一维序列中,并传递给一个类似ViT的骨干来提取表示。SwinUNet在Swin Transformer中的Transformer块的基础上建立了一个基于U型Transformer的分割模型,其中取得了显著的改进。DS-TransUNet进一步扩展了SwinUNet,增加了一个编码器来处理多尺度输入,并引入了一个融合模块,通过自注意力机制有效地建立不同尺度的特征之间的全局依赖关系。与这些基于Transformer的主干相比,nnFormer继承了卷积在编码精确的空间信息和产生分层表示方面的优势,有助于在不同尺度上建立物体模型概念。
nnFormer的整体结构如图2所示,它保持了与U-Net类似的U型结构,主要由三部分组成,即编码器、瓶颈和解码器。具体来说,编码器包括一个嵌入层、两个局部Transformer块(每个块包含两个连续的层)和两个下采样层。对称的是,解码器分支包括两个Transformer块,两个上采样层和最后一个用于进行Mask预测的patch扩展层。此外,瓶颈包括一个下采样层、一个上采样层和三个全局Transformer块,用于提供大的感受野以支持解码器。受U-Net的启发,我们在编码器和解码器的相应特征金字塔之间以对称的方式添加了跳过连接,这有助于恢复预测中的细粒度细节。然而,与通常使用求和或串联操作的非典型跳过连接不同,我们引入了跳过关注来弥补编码器和解码器之间的差距。
在下文中,我们将在Synapse数据集上演示前向程序。根据Synapse上的程序,可以很容易地推断出不同数据集上的结果。
nnFormer的输入是一个三维patch X∈RH×W×D(通常是从原始图像中随机裁剪的),其中H、W和D分别表示每个输入扫描的高度、宽度和深度。
在Synapse上,嵌入块负责将每个输入扫描X转化为高维张量Xe∈RH/4×W/4×D/2×C,其中H/4×W/4×D/2代表patch tokens的数量,C代表序列长度(这些数字在不同的数据集上可能略有不同)。与ViT和Swin Transformer在embedding block中使用大的卷积核来提取特征不同,我们发现应用小的卷积核的连续卷积层在初始阶段带来更多的好处,这可以从两个角度来解释,即:i)为什么应用连续的卷积层和ii)为什么使用小尺寸核。对于i),我们在embedding block中使用卷积层,因为它们对像素级的空间信息进行编码,比变换器中使用的patch式位置编码更精确。对于ii),与大尺寸的内核相比,小的内核尺寸有助于降低计算的复杂性,同时提供同等大小的感受野。如图2b所示,嵌入块由四个卷积层组成,其核大小为3。在每个卷积层之后(除了最后一个),附加一个GELU和一个LN层。在实践中,根据输入patch的大小,嵌入块中的卷积层的步长可能会相应变化。
在embedding block之后,我们将高维张量Xe传递给Transformer块。背后的要点是将捕捉到的长期依赖关系与下采样层产生的不同尺度的分层对象信息和初始嵌入层编码的高分辨率空间信息完全融合在一起。与Swin Transformer相比,我们在三维局部体积(即LV-MSA)内计算自注意力,而不是二维局部窗口。
假设XLV∈RL×C代表局部Transformer块的输入,XLV将首先被重塑为XLV帽∈RNLV×NT×C,其中NLV是预先定义的三维局部体积的数量,NT = SH × SW× SD表示每个体积中patch标记的数量。{SH ,SW,SD}代表局部体积的大小。
如图3a所示,我们遵循Swin Transformer在每个块中进行两个连续的变换层,其中第二层可以被看作是第一层的移位版本(即SLV-MSA)。主要的区别在于,我们的计算是建立在三维局部体积之上,而不是二维局部窗口。
计算过程可以总结为以下几点。
这里,l代表层的索引。MLP是多层感知器(multi-layer perceptron)的缩写。LV-MSA在一个h×w×d大小patch的体积上的计算复杂度为:
其中Q、K、V∈RNT×dk表示query、key和value矩阵。B∈RNT是相对位置编码。在NT = SH × SW× SD表示每个体积中patch标记的数量。实践中,我们首先初始化一个较小尺寸的位置矩阵ˆB∈R(2SH-1)×(2SW-1)×(2SD-1),并从ˆB中取相应的值来建立一个较大的位置矩阵B。
我们发现,通过将Swin中的patch合并操作替换为直接的分层卷积,nnFormer可以对体积图像分割提供更多的改进。其背后的直觉是卷积下采样产生分层表示,有助于在多个尺度上对对象进行建模。
如图2c所示,在大多数情况下,下采样层涉及一个跨度卷积操作,其中跨度在所有维度上都被设置为2。然而,在实践中,特定维度的步长可以设置为1,因为在这个维度上,切片的数量是有限的,过度下采样(即使用大的下采样步长)可能是有害的
最初的ViT采用了原始的二维多头自我注意机制。
在本文中,我们将其扩展到三维版本(如图3b所示),其计算复杂度可以表述为:
(4)与(2)相比,很明显,当{h, w, d}比{SH, SW , SD}相对较大(例如,大一个数量级)时,GV-MSA需要更多的计算资源。事实上,这正是我们在编码器中使用局部Transformer块的原因,这些Transformer块的设计是为了通过局部自注意力机制有效地处理大尺寸的输入。
然而,在瓶颈部分,{h,w,d}在经过几个下采样层后已经变得非常小,使得它们的乘积,即hwd,具有与SH SW SD相似的大小。
这就为应用GV-MSA创造了条件,与LV-MSA相比,GV-MSA能够提供更大的感受野,大感受野在不同的应用中被证明是有益的。在实践中,我们在瓶颈处使用了三个全局Transformer块(即六个GV-MSA层)来为解码器提供足够的感受野信息。
解码器中的两个Transformer块的结构与编码器中的Transformer块是高度对称的。与下采样块相比,我们采用分层的反卷积将低分辨率的特征图上采样为高分辨率的特征图,而这些特征图又通过跳过注意力与来自编码器的表征合并,以捕捉语义和细粒度的信息。与上采样区块类似,最后一个patch扩展区块也采取反卷积操作来产生最终的Mask预测。d图是上采样层内部结构,e图是扩展层的
跳过注意力。 convnets[Unet, Resnet]中的非典型跳过连接适应串联或求和,以纳入更多信息。受Transformer中机器翻译任务的启发,我们建议用一种注意力机制来取代串联/求和,本文将其命名为跳过注意力。具体来说,编码器的第l个Transformer块的输出,即Xl{LV,GV},经过线性投影(即单层神经网络)后,被转换并分割成一个key矩阵Kl∗和一个value矩阵Vl∗。
LP是线性投影层之后的,利用编码器未下采样之前与,经过解码器上采样之后的值进行计算,编码器投射后的作为K和V,而解码器上采样后的作为Q。
为了彻底比较nnFormer和以前基于convnet和transformer的架构,我们在三个数据集/任务上进行了实验:医学分割十项全能(MSD)中的脑瘤分割任务、Synapse多器官分割和自动心脏诊断挑战赛(ACDC)。对于每个实验,我们重复十次,并报告其平均结果。我们还计算了p-values,以证明nnFormer的意义。
使用MRI扫描进行脑肿瘤分割
这项任务包括484张MRI图像,每张图像包括四个通道,即FLAIR、T1w、T1gd和T2w。这些数据来自19个不同的机构,包含2016年和2017年脑肿瘤分割(BraTS)挑战中使用的数据子集。相应的目标 ROI是三个肿瘤子区域,即水肿(ED)、增强型肿瘤(ET)和非增强型肿瘤(NET)。为了与UNETR中报道的那些结果一致,在比较我们的nnFormer和基于Transformer的模型时,我们显示了整个肿瘤(WT)、增强型肿瘤(ET)和肿瘤核心(TC)的实验结果。对于数据的分割,我们遵循UNETR的指示,其中训练/验证/测试集的比例分别为80%、15%和5%。如上所述,我们使用HD95和Dice分数作为评价指标。
用于多器官CT分割的Synapse
这个数据集包括30例腹部CT扫描。按照[TransUNet]中使用的分割方法,提取18个病例建立训练集,其余12个病例用于测试。我们报告了用95% Hausdorff距离(HD95)和Dice评分(DSC)对8个腹部器官进行评估的模型性能,这些器官是主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺和胃。
ACDC用于自动化心脏诊断
ACDC涉及100个病人,右心室的腔体、左心室的心肌和左心室的腔体要进行分割。每个病例的标签涉及左心室(L V)、右心室(RV)和心肌(MYO)。
数据集被分成70个训练样本,10个验证样本和20个测试样本。评价指标包括HD95和Dice分数
我们基于Python 3.6、PyTorch 1.8.1和Ubuntu 18.04运行所有实验。所有的训练程序都是在一个拥有11GB内存的NVIDIA 2080 GPU上进行的。初始学习率被设定为0.01,我们采用了方程7中描述的 "poly "衰减策略。默认的优化器是SGD,我们将动量设置为0.99。权重衰减被设置为3e-5。我们利用交叉熵损失和Dice损失,简单地将它们相加。训练epoch的数量(即公式7中的最大epoch)为1000,一个epoch包含250次迭代。在Synapse上,不同编码器阶段使用的多头自我注意的头数为[6, 12, 24, 48]。在其余两个数据集中,头数变为[3、6、12、24]。
**预处理和增强策略。**所有的图像将首先被重新取样到相同的目标间距。在训练过程中,旋转、缩放、高斯噪声、高斯模糊、亮度和对比度调整、模拟低分辨率、伽马增强和镜像等增强措施按给定顺序应用。
**深度监督。**我们还在训练阶段增加了深度监督。具体来说,解码器中每个阶段的输出被传递到最后的扩展块,在那里将应用交叉熵损失和Dice损失。在实践中,考虑到一个典型阶段的预测,我们对Ground Truth的Mask进行下采样,以匹配预测的分辨率。因此,最终的训练目标函数是三个分辨率下所有损失的总和。
网络配置。 在表I中,我们显示了所有三个数据集上的实验的网络配置。与nnUNet相比,在nnFormer中,较小尺寸的输入patch可以获得更好的分割结果。
**脑肿瘤分割。**表二列出了所有模型在脑瘤分割任务上的实验结果。我们的nnFormer在所有类别中取得了最低的HD95和最高的DSC分数。此外,nnFormer能够在相同的评估指标上以较大的优势超过第二好UNETR的方法
例如,nnFormer在HD95的平均值中比UNETR高出4.5毫米,在DSC的每个类别中高出近10个百分点。与以前基于Transformer的方法相比,nnFormer在HD95中显示出比DSC更强的实力。
多器官分割(Synapse)。如表三所示,我们对Synapse进行了实验,并将我们的nnFormer与各种基于Transformer的方法进行比较。
我们可以看到,表现最好的方法是LeViT-UNet-384s和TransUNet。LeViT-UNet-384s实现了平均16.84毫米的HD95,而TransUNet产生了平均DSC为84.36%的结果。相比之下,我们的nnFormer在平均HD95和DSC方面能够超过LeViT-UNet-384s和TransUNet,分别超过6毫米和2个百分点,这是对Synapse相当了不起的改进。具体来说,nnFormer在六个器官中实现了最高的DSC,包括Aotra、肾脏(左)、肾脏(右)、肝脏、胰腺和胃。与以前基于Transformer的方法相比,nnFormer在分割胰腺和胃方面更有优势,这两个器官用过去的分割模型是很难划分的。
自动心脏诊断(ACDC)。表四显示了ACDC的实验结果。
我们可以看到,最好的基于Transformer的模型是LeViT-UNet-384s,其平均DSC略高于SwinUNet,而TransUNet和SwinUNet在处理左心室(LV)的划分方面更有能力。相比之下,nnFormer在所有类别中都超过了LeViT-UNet-384s,并且在平均DSC上超过了近1.7个百分点,这再次验证了它比过去基于Transformer的方法的优势。
**统计学意义。**在表II、III和IV中,我们采用独立的双样本t检验来计算我们的nnFormer的平均性能与HD95和DSC中表现最好的基线之间的p值。无效假设是,我们的nnFormer与表现最好的基线相比没有优势。我们可以看到,在所有三个公共数据集上,nnFormer在HD95和DSC下产生的p值都小于1e-2,这表明反对无效假设的有力证据。因此,nnFormer在三个不同的任务上显示出比以前的基于Transformer的方法有明显的改进。
在这一部分,我们将nnFormer与nnUNet进行比较,后者被认为是最强大的三维医学图像分割模型之一。
**结果。**在表五中,我们显示了HD95和DSC指标的具体等级结果,以进行全面的比较。
具体来说,从特定类别的HD95结果来看,nnFormer在16个类别中的11个方面优于nnUNet。在特定类别的DSC中,nnFormer在16个类别中的9个方面优于nnUNet。因此,在HD95下,nnFormer似乎更有优势,这意味着nnFormer可能更好地划定对象边界。从平均性能来看,我们可以看到nnFormer往往能取得更好的平均性能。例如,在HD95结果较低的三个公共数据集上,nnFormer的表现优于nnUNet,而在DSC结果较高的三个数据集中,有两个数据集的表现优于nnUNet。
**统计学意义。**为了进一步验证nnFormer对nnUNet的意义,我们还计算了nnFormer和nnUNet的平均性能之间的p值。
与我们在表二中所做的类似,我们分别提供了基于HD95和DSC的三个公共数据集的两个p值。最明显的观察是,nnFormer在HD95的三个公共数据集上取得的p值小于0.05。这些结果表明,当HD95被作为主要评价指标时,nnFormer是首选。此外,基于DSC的肿瘤和多器官分割的p值(>0.05)意味着nnFormer是一个与nnUNet相当的模型,而ACDC上的结果证明了nnFormer的重要性。总之,在DSC下,nnF ormer比nnUNet略有优势。
模型组合。除了单个模型的性能,我们还研究了nnFormer和nnUNet之间的多样性,这是模型集合的一个关键因素。令人惊讶的是,我们发现,通过简单地对nnFormer和nnUNet的预测进行平均化(即表五中的nnAvg),已经可以将整体性能提高很大幅度。例如,nnAvg在HD95和DSC下的肿瘤分割中取得了所有类别的最佳结果。此外,当评价指标为HD95时,nnAvg在Synapse提高了近30%。这些结果表明,nnFormer和nnUNet是高度互补的。
表六显示了我们对nnFormer中不同模块的消融研究结果。为了简单起见,我们在ACDC上进行了实验,并使用DSC作为默认的评估指标。
表六中最基本的基线(第0行)包括LV-MSA(但没有SLV -MSA),在Swin中使用的patch Merging和embed layer。我们可以看到,这样的组合已经可以达到比LeViT-UNet-38更高的平均DSC,后者是表四中表现最好的基线。
【1】我们首先用我们提出的体积嵌入层,即小卷积核大小和卷积步长的连续卷积层,取代了以大的卷积核大小和卷积步长实现的patch嵌入层。我们发现,引入的卷积嵌入层将平均DSC提高了约0.4个百分点。
【2】接下来,我们去掉了patch Merging层,加入了我们的卷积下采样层。我们发现这种简单的替换可以进一步提高整体性能0.3个百分点。
【3】然后,我们用GV-MSA替换了瓶颈处的LV-MSA,我们观察到0.2%的改进。这一现象表明,提供足够大的感受野对分割任务是有益的。
【4】之后,我们用跳过注意力来代替传统的连接/求和操作。令人惊讶的是,我们发现跳过注意力能够将整体性能提高0.4个百分点,这表明跳过注意力可以作为传统跳过连接之外的另一种选择。
【5】最后但并非最不重要的是,我们研究了在每个Transformer块中增加更多的Transformer层,方法是像Swin Transformer中那样将一个SLV-MSA层与每个LV-MSA层级联,并将全局自注意力层的数量增加一倍。我们发现,引入更多的Transformer层确实给整体性能带来了更多的改善,因为它将更多的长程依赖关系混合在所学的体积表征中。
在图4中,我们可视化了我们的nnFormer、nnUNet和UNETR在三个公共数据集上的一些分割结果。
与UNETR相比,我们的nnFormer可以大大减少假阳性预测的数量。一个典型的例子是ACDC上的第五个例子。我们可以看到,UNETR在心肌外产生了大量错误的右心室像素。相比之下,我们的nnFormer没有产生心肌外的右心室预测,这表明nnFormer在ACDC上比UNETR更有鉴别力。
另一方面,我们观察到nnUNet显示出非常有竞争力的分割结果,在几乎所有的例子中都比UNETR好很多。然而,我们仍然发现nnFormer与nnUNet相比保持着明显的优势,其中之一就是nnFormer更善于处理边界问题。事实上,这一现象已经反映在表六中,当HD95为默认评价指标时,nnFormer明显优于nnUNet。
在本文中,我们提出了一种用于3D图像分割的三维Transformer:nnFormer。nnFormer是在卷积和自注意力的交错主干之上构建的。
卷积有助于编码精确的空间信息和构建层次分明的对象概念。对于自注意力,nnFormer采用了三种注意机制来纠缠长距离的依赖关系。具体来说,基于局部和全局体积的自注意力侧重于构建特征金字塔和提供大的感受野。跳过式注意力负责弥补编码器和解码器之间的差距。实验表明,在HD95和DSC中,nnFormer与以前基于Transformer的模型相比,保持了很大的优势。与nnUNet相比,nnFormer在HD95中的表现明显更好,同时在DSC中产生了相当的结果。更重要的是,我们证明了nnFormer和nnUNet在模型组合中可以相互受益,其中简单的平均化操作已经可以产生很大的改进。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。