nnFormer
nnFormer不仅利用了交错卷积和自注意操作的结合,还引入了局部和全局基于体积的自注意机制来学习体积表示。此外,nnFormer还提出在类似U-Net的架构中,使用跳转注意(skip-attention)来取代跳转连接(skip-connection)中传统的连接/求和操作.
有许多方法试图将变压器应用于医学图像分析领域。Chen等人[11]首次提出TransUNet来探索变形器在医学图像分割中的潜力。TransUNet的整体架构类似于U-Net[12],其中convnets充当特征提取器和转换器,帮助编码全局上下文。事实上,TransUNet及其大多数追随者[13]-[16]的一个主要特点是将convnets视为主体,并在其基础上进一步应用变压器来捕获长期依赖关系。然而,这样的特点可能会引起一个问题,那就是变压器的优点没有得到充分利用。换句话说,我们认为一层或两层转换器不足以将长期依赖关系与卷积表示纠缠在一起,卷积表示通常包含精确的空间信息并提供分层概念。
Karimi等人[17]首次引入了一种无卷积的分割模型,将扁平的图像表示转发给变压器,然后将变压器的输出重组为三维张量,与分割蒙版对齐。最近,Swin Transformer[3]表明,通过参考convnets中使用的特征金字塔,变压器可以通过对特征映射进行适当的下采样来学习不同尺度的分层对象概念。受此想法的启发,SwinUNet[18]利用分层变压器块在类似U-Net的架构中构建编码器和解码器,DS-TransUNet[19]在此基础上增加了一个编码器以接受不同大小的输入。SwinUNet和DS-TransUNet都实现了对TransUNet的持续改进。然而,他们没有探索如何将卷积和自关注适当地结合起来构建最优的医疗分割网络。
各组件的功能
首先,我们在变压器块前面放置了一个轻量级的卷积嵌入层。与[17]中直接平坦化原始像素并进行1D预处理相比,卷积嵌入层编码精确(即像素级)的空间信息,并提供低级但高分辨率的3D特征。在嵌入块之后,将变压器和卷积下采样块相互交错,在不同尺度上将长期依赖关系与高级和分层的对象概念充分地交织在一起,有助于提高学习表征的泛化能力和鲁棒性。
nnFormer的另一个贡献在于提出了一种计算效率高的方法来利用片间依赖关系。具体而言,nnFormer提出联合使用Local Volume-based Multi-head Self-attention (LV-MSA)和Global Volume-based Multi-head Self-attention (GV-MSA)构建特征金字塔,为学习局部和全局三维体上的表征提供足够的接受场,然后将其聚合进行预测。与朴素多头自注意(MSA)方法[1]相比,该方法在产生具有竞争力的分割性能的同时,大大降低了计算复杂度。此外,受机器翻译任务中使用的注意机制[1]的启发,我们引入了跳过注意来取代类似U-Net架构的跳过连接中的非典型拼接/求和操作,进一步改善了分割结果。
关于transformer在医学图像分割上的应用:
1、基于卷积的
TransUNet[11]首次应用变压器改进医学图像的分割结果。TransUNet将convnet作为特征提取器来为输入片生成特征映射。然后将补丁嵌入应用于瓶颈中的特征映射的补丁,而不是ViT中的原始图像[2]。同时,与TransUNet类似,Li等人[20]提出使用压缩注意块来正则化变形器的自注意模块,使用扩展块来学习眼底图像的多样化表示,这些都是在convnets内部的瓶颈中实现的。
TransFuse[13]引入了一个BiFusion模块,将基于浅convnet的编码器和基于变压器的分割网络的特征融合在一起,对2D图像进行最终预测。与TransUNet相比,TransFuse主要将自关注机制应用于输入嵌入层,以改进二维图像的分割模型。Yun等[21]采用变压器将光谱信息与卷积特征编码的光谱信息纠缠在一起,以解决高光谱病理问题。
Xu等人[22]广泛研究了变压器和convnets之间的权衡,并提出了一种更高效的编码器,名为levi - unet。Li等人[23]提出了一种新的上采样方法,并将其纳入UNet的解码器中,对长期依赖关系和全局信息进行建模,以获得更好的重建结果。
TransClaw U-Net[15]在具有更多卷积特征金字塔的UNet中使用了变压器。
TransAttUNet[16]探讨了将变压器自注意与卷积全局空间注意相结合的可行性。
Xie等[24]采用变压器从不同的convnets层捕获多尺度卷积特征的长期依赖关系。
TransBTS[25]首先利用3D convnets提取体积空间特征,并对输入的3D图像进行下采样以产生分层表示。TransBTS中编码器的输出然后被重塑成一个向量(即token),并被馈送到变压器中进行全局特征建模,之后,一个普通的卷积解码器被附加到上样特征映射中以实现重建的目标。与这些直接使用卷积神经网络作为特征提取器的方法不同,我们的nnFormer在功能上依赖于卷积和基于变换的块,它们相互交错以利用彼此的优势。
2、基于transformer的
Valanarasu等人[14]提出了一种门控轴向注意模型(即MedT),该模型通过在自注意中引入汇总控制机制,扩展了现有的convnet架构。Karimi等[17]去除了卷积运算,建立了基于变压器的三维分割模型。主要思想就是将本地体积块分成patch。然后将其平面化并嵌入到1D序列中,并传递到类似于vita的主干以提取表示。SwinUNet[18]在[3]中的变压器块的基础上构建了基于u形变压器的分割模型,并取得了明显的改进。DS-TransUNet[19]进一步扩展了SwinUNet,增加了一个编码器来处理多尺度输入,并引入融合模块,通过自关注机制有效建立不同尺度特征之间的全局依赖关系。与这些基于变换的系统相比,nnFormer继承了卷积在编码精确空间信息和生成分层表示方面的优势,有助于在各种尺度上对对象概念进行建模。
细分来看:
1、Encoder:
nnformer的输入是一个3D patch X∈RH×W ×D(通常从原始图像中随机裁剪),其中h、W、D分别表示每次输入扫描的高度、宽度和深度。
The embedding layer.
在Synapse上,嵌入块负责将每个输入扫描X转换为高维张量Xe∈R H/4 × W/4 × D/2 ×C,其中H/4 × W/4 × D/2表示补丁标记的数量,C表示序列长度(这些数字在不同的数据集上可能略有不同)。与ViT[2]和Swin Transformer[3]在嵌入块中使用大的卷积核来提取特征不同,我们发现在初始阶段使用连续的卷积层和小的卷积核可以带来更多的好处,这可以从两个角度来解释,即i)为什么要使用连续的卷积层,ii)为什么要使用小的内核。对于i),我们在嵌入块中使用卷积层,因为它们编码像素级空间信息,比变压器中使用的patch-wise位置编码更精确。对于ii),与大内核相比,小内核尺寸有助于降低计算复杂度,同时提供相同大小的接受域。如图2b所示,嵌入块由四个卷积层组成,每个卷积层的内核大小为3。在每个卷积层(最后一个除外)之后,追加一个GELU[26]层和一个层归一化[27]层。在实际应用中,根据输入patch的大小,嵌入块中的卷积步幅也会相应变化。(这里应该是用多个小卷积核来取代大的卷积核)
Local Volume-based Multi-head Self-attention (LV-MSA)
在嵌入层之后,我们将高维张量Xe传递给变压器块。其主要目的是将捕获的长期依赖关系与下采样层产生的各种尺度的分层对象概念和初始嵌入层编码的高分辨率空间信息完全纠缠在一起。与Swin Transformer[3]相比,我们在3D局部体积(即LV-MSA,基于局部体积的多头自关注)内计算自关注,而不是在2D局部窗口内计算。设XLV∈RL×C表示局部变压器块的输入,则XLV首先被重构为XLV∈RNLV×NT ×C,其中NLV为三维局部卷的预定义数量,NT = SH × SW × SD为每个卷中的补丁令牌数量。{SH, SW, SD}表示本地卷的大小。如图3a所示,我们按照[3]在每个区块中连续进行两层变压器,其中第二层可以看作是第一层的移位版本(即SLV-MSA)。主要的区别在于我们的计算是建立在3D局部体积之上,而不是2D局部窗口。计算过程可总结如下:
复杂度为
SLV-MSA将LV-MSA替换成,
,其中,B是相对位置编码,我们首先初始化一个较小尺寸的位置矩阵,并从中取相应的值来构建一个较大的位置矩阵B。
2、The down-sampling layer
我们发现用直接的strided卷积(跨步卷积)代替[3]中的补丁合并操作,nnFormer可以在体积图像分割方面提供更多的改进,这是因为卷积下采样产生层次表示,有助于在多个尺度上建模对象概念。如图2c所示,在大多数情况下,下采样层涉及一个跨步卷积操作,其中所有维度的步幅都设置为2。然而,在实践中,可以将特定维度的步幅设置为1,因为该维度的切片数量有限,过度向下采样(即使用较大的向下采样步幅)可能是有害的。
3、Bottleneck
vit采用了原始的2D多头自注意机制,在这里,我们将其扩展为3D版本,其复杂度为
与(2)相比,很明显,当{h, w, d}相对于{SH, SW, SD}较大(例如,大一个数量级)时,GV-MSA需要更多的计算资源。事实上,这正是我们在编码器中使用局部变压器块的原因,其设计目的是通过局部自关注机制有效地处理大尺寸的输入。然而,在瓶颈中,{h, w, d}经过几层下采样后已经变得非常小,使得它们的乘积hwd,具有与SH SW SD相似的大小。这就为GV-MSA的应用创造了条件,GV-MSA比LV-MSA能够提供更大的感受野,而大的感受野已被证明在不同的应用中是有益的[28]-[31]。在实践中,我们在瓶颈中使用三个全局变压器块(即六个GV-MSA层)来为解码器提供足够的接受域。
4、Decoder
解码器中的两个变压器块的结构与编码器中的变压器块高度对称。与下采样块相比,我们采用跨步反卷积将低分辨率特征映射上采样到高分辨率特征映射,然后通过跳过注意与编码器的表示合并,以捕获语义和细粒度信息。与上采样块类似,最后一个补丁扩展块也需要反卷积操作来产生最终的掩码预测。
5、Skip Attention
convnets中的非典型skip-connection[12,32]采用串联或求和来包含更多信息。受文献[1]中机器翻译任务的启发,我们提出用一种注意机制来代替连接/求和,本文将其命名为skip-attention。具体来说,编码器的第l变压器块的输出,即,经过线性投影(即单层神经网络)后,变换并分裂为key矩阵Kl∗和value矩阵V l∗
其中LP为线性投影。因此,在解码器的第l∗个上采样层之后的输出特征映射X l∗UP被视为Query Ql∗。然后,我们可以像在(3)中所做的那样,对解码器中的Ql∗,Kl∗和V l∗进行LV/GV-MSA,即:
其中l *表示层索引。dl * k和Bl *与(3)中的意义相同,因此它们的大小可以很容易地推断出来。