【图像分割 2023 WACV】HiFormer
论文题目：HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation
中文题目：HiFormer:基于Transformer的分层多尺度表示医学图像分割
论文链接：
论文代码：
论文团队：
发表时间：
DOI:
引用：
引用数：
摘要
卷积神经网络(cnn)已成为医学图像分割的共识。然而，由于卷积运算的性质，它们在模拟远程依赖关系和空间相关性方面受到限制。虽然Trasnformer最初是为了解决这个问题而开发的，但它们无法捕获低级特征。相比之下，研究表明，局部和全局特征对于密集预测至关重要，例如在具有挑战性的环境中进行分割。在本文中，我们提出了一种新的方法HiFormer，它有效地桥接了CNN和Trasnformer，用于医学图像分割。具体来说，我们使用开创性的Swin Transformer模块和基于cnn的编码器设计了两个多尺度特征表示。为了确保从上述两种表示中获得的全局和局部特征的良好融合，我们在编码器-解码器结构的跳过连接中提出了双级融合(DLF)模块。在各种医学图像分割数据集上的大量实验表明，HiFormer在计算复杂度、定量和定性结果方面优于其他基于cnn、基于Trasnformer和混合方法。我们的代码在GitHub上是公开的。
1. 介绍
医学图像分割是计算机视觉的主要挑战之一，它提供了详细的解剖学领域所需的有价值的信息分析。这些信息可以极大地帮助医生描述损伤，监测疾病进展，并评估是否需要适当的治疗。由于越来越多地使用医学图像分析，高度精确和鲁棒的分割变得越来越重要。
卷积神经网络(cnn)具有令人印象深刻的图像特征提取能力，已广泛用于不同的图像分割任务。随着基于编码器-解码器的网络的兴起，如全卷积网络[36]，u形结构，如:U-Net[41]及其变体，cnn在医学图像分割任务中取得了显著的成功。在这两种结构中，跳过连接分别用于体现编码器和解码器路径提供的高级和细粒度特性。尽管CNN模型在各种计算机视觉任务中取得了成功，但由于其有限的接受野和固有的归纳偏差，其性能受到限制[19,4]。上述原因阻止了cnn在图像中构建全局上下文和长期依赖关系，从而限制了它们在图像分割中的性能。
最近，受自然语言处理(NLP)中Trasnformer的杰出成功[47]的激励，视觉Trasnformer被开发出来以减轻cnn在图像识别任务中的缺点[19]。transformer主要利用多头自关注(MSA)机制，该机制可以有效地在令牌序列之间构建远程依赖关系并捕获全局上下文。vanilla vision transformer[19]表现出与基于cnn的方法相当的性能，但需要大量数据进行泛化，并且具有二次复杂度。有几种方法被提议解决这些限制。DeiT[44]提出了一种高效的知识蒸馏训练方案，克服了视觉Trasnformer需要大量数据学习的困难。Swin Transformer[35]和pyramid vision Transformer[48]分别尝试利用基于窗口和空间约简的注意力来降低视觉转换器的计算复杂度。
此外，近年来，多尺度特征表示在视觉变换中表现出了强大的性能。CrossViT[12]提出了一种新的双支路Trasnformer架构，提取多尺度上下文信息，为图像分类提供更细粒度的特征表示。同样，DS-TransUNet[33]提出了一种双支路Swin-Transformer，在编码器中捕获不同的语义尺度信息，用于医学图像分割任务。HRViT[23]将多分支高分辨率架构与视觉转换器连接起来进行语义分割。因此，这样的结构可以有效地帮助增强令牌之间的远程关系的建模，并获得更详细的信息。
尽管视觉Trasnformer有能力对全局上下文表示进行建模，但自注意机制导致了低层次特征的缺失。为了缓解上述问题，已经提出了混合cnn-transformer方法，利用cnn的局域性和transformer的远程依赖特性对全局和局部特征进行编码，特别是医学图像分割中的TransUnet[13]和LeVit-Unet[52]。然而，这些方法存在一些阻碍它们获得更高性能的障碍:1)它们不能在保持特征一致性的同时有效地结合低级和高级特征;2)它们不能很好地利用分层编码器产生的多尺度信息。
在本文中，我们提出了一种新的基于cnn-transformer的编码器-解码器框架，该框架有效地利用了transformer的全局远程关系和cnn的局部特征表示来完成精确的医学图像分割任务。编码器由三个模块组成:两个分层CNN和SwinTrasnformer模块和DLF模块。Swin Transformer和CNN模块各包含三个级别。首先，将输入图像送入CNN模块以学习其局部语义表示。为了弥补全局表示的不足，Swin Transformer模块被应用于CNN的浅层特征之上，以捕获远程依赖关系。接下来，利用具有不同窗口大小的Swin Transformer模块的金字塔来学习多尺度交互。为了鼓励特征可重用性并提供本地化信息，设计了一个跳过连接模块来将CNN的本地特征转移到Transformer块中。然后将最小和最大金字塔层的最终表示形式输入DLF模块。提出的DLF模块是一个多尺度视觉转换器，利用交叉注意机制融合两个获得的特征映射。
最后，将两个重新校准的特征映射传递到解码器块中以产生最终的分割掩码。我们提出的HiFormer不仅缓解了上述问题，而且在不同的评价指标上都超过了所有同类产品。我们的主要贡献:•一种新的混合方法，融合了Trasnformer的远程上下文交互和CNN的局部语义信息。
•DLF模块在粗粒度和细粒度特征表示之间建立有效的特征融合。
实验结果证明了该方法在医学图像分割数据集上的有效性和优越性。
2. 相关工作
2.1 基于cnn的分割网络
卷积神经网络被认为是不同计算机视觉任务的事实上的标准。cnn取得优异成绩的一个领域是图像分割，其中为每个像素分配类标签。Long等[36]表明，完全卷积网络(fully convolutional networks, fcn)可以用于分割没有完全连接层的图像。考虑到卷积层按顺序堆叠的普通fcn的输出通常是粗糙的，因此提出了融合不同层输出的其他模型[6,39,41]。已经引入了几种方法来改善FCN的有限接受野，包括扩展卷积[14,54]和上下文建模[56,15]。CNN模型在医学成像任务中表现出色。在引入U-net后[41]，其他研究人员将重点放在利用u型编码器-解码器结构上。在[46]中，U-net扩展了一个过完备的网络，而在Unet++[57]中，编码器-解码器架构被重新设计，在模块之间增加了密集的跳过连接。该结构已被进一步改进并应用于不同的医学领域[8,30,21,5]。
2.2 VIT
继Trasnformer在NLP中取得显著成功后[47]，Dosovitskiy等[19]提出了视觉Trasnformer(Vision Transformer, ViT)，该方法利用自注意机制学习全局信息，在图像分类任务上取得了最先进的性能。为了提高视觉转换器的效率，减少对大型数据集的依赖，以实现泛化，已经引入了几种视觉转换器的衍生产品[44,55,51]。
此外，人们还提出了许多方法，重点是多尺度表示，通过从不同尺度提取信息来提高准确性和效率。灵感来自cnn的金字塔结构[38,53,9,34]， PVT[48]是第一个引入的金字塔视觉Trasnformer。后来，Swin Transformer[35]提出了一种分层视觉转换器，使用有效的移位窗口方法来局部计算自注意力。
CrossViT[12]建议在线性时间内使用双分支视觉转换器，然后使用交叉注意模块，以获得更丰富的特征表示。视觉转换器在其他视觉任务中也显示出令人印象深刻的结果，包括[58,20]，它提供基于端到端转换器的对象检测模型，以及[43,24]的语义和实例分割。
2.3 用于医学图像分割的Transformer
尽管CNN模型取得了令人鼓舞的结果，但由于其有限的接受域，这种方法通常在建模远程依赖关系时表现出限制，从而产生较弱的性能。最近，在医学图像分割中，基于变换的模型比CNN模型更受欢迎。Swin- unet[10]和DS-TransUNet[33]提出了基于Swin transformer的u型结构纯Trasnformer模型，用于二维分割。除了全Trasnformer模型外，TransUNet[13]还利用cnn和Trasnformer来捕获低级和高级特征。UNETR[27]使用基于Trasnformer的编码器嵌入输入的3D patch，并使用基于cnn的解码器实现最终的3D分割结果。大多数先前的工作要么使用缺乏全局特征的CNN，要么使用局限于局部特征表示的Trasnformer进行特征提取;这将导致不包含丰富信息的无效特性映射。在混合作品中，采用简单的特征融合机制，不能保证不同尺度之间的特征一致性。在多尺度表示的激励下，我们提出了一种基于cnnTrasnformer的结构HiFormer，它有效地融合了全局和局部信息，并利用一种新的基于Trasnformer的融合方案来保持二维医学图像分割任务的特征丰富性和一致性。
3. 方法
本节将概述拟议的HiFormer。如图1a所示，我们提出的架构提供了一个端到端的训练策略，该策略集成了Swin Transformer的全局上下文表示和编码器中CNN模块的本地代表性特征。然后使用双级融合模块(DLF)获得更丰富的特征表示。之后，解码器输出最终的分割图
3.1 编码器
如图1a所示，所提出的编码器由CNN和Swin Transformer两个层次模型组成，DLF模块丰富检索到的特征，并准备将其馈送到解码器中。由于单独使用CNN或Trasnformer会导致局部或全局特征被忽略，从而影响模型的性能，因此我们首先利用CNN局域性特征来获取局部特征。在这里，CNN和Swin Transformer都包括三个不同的级别。我们通过跳过连接将每一层的局部特征转移到相应的Swin Transformer层，以获得通用表示。然后将每个转移的CNN电平与其并联的Trasnformer电平相加，并通过Patch merged模块生成分层表示(见图1a)。我们利用分层设计来利用多尺度表示。最大和最小的级别进入DLF模块，以交换来自不同尺度的信息并生成更强大的功能。
在下面，我们将深入详细地讨论我们的CNN, SwinTrasnformer和DLF模块。
3.1.1 CNN模块
提出的编码器首先采用CNN作为特征提取器，构建不同分辨率的中间CNN特征图的金字塔。取一个空间维度为H、W、通道为C的输入图像X∈RH×W×C，首先将其送入CNN模块。CNN模块由三层组成，通过1 × 1的Conv将跳跃式连接连接到相关联的Trasnformer层，补偿Trasnformer的底层缺失信息，恢复局部空间信息。
3.1.2 SwinTransformer 模块
普通的Trasnformer编码器块[19]由两个主要模块组成:多头自注意(MSA)和多层感知器(MLP)。香草Trasnformer由N个相同的Trasnformer编码器块组成。在每个块中，在MSA和MLP块之前，应用LayerNorm (LN)。此外，激活的副本还通过跳过连接添加到MSA或MLP块的输出中。使用标准MSA的香草ViT的一个主要问题是它的二次复杂度，使得它在执行图像分割等高分辨率计算机视觉任务时效率低下。为了克服这一限制，Swin Transformer[10]引入了WMSA和SW-MSA。
Swin Transformer模块包括两个连续修改的Trasnformer模块;MSA块被基于窗口的多头自注意(W-MSA)和移位窗口的多头自注意(SW-MSA)所取代。在W-MSA模块中，自关注将应用于大小为 M × M M \times M M×M的局部窗口。然而，考虑到没有跨窗口的连接，它的建模能力有限。为了缓解这一问题，引入了SW-MSA，它利用与W-MSA模块的输入相比移位的窗口配置;这是为了确保我们有跨窗口连接。描述了这个过程
z ^ l = W − M S A ( L N ( z l − 1 ) ) + z l − 1 , z l = M L P ( L N ( z ^ l ) ) + z ^ l , z ^ l + 1 = S W − M S A ( L N ( z l ) ) + z l , z l + 1 = M L P ( L N ( z ^ l + 1 ) ) + z ^ l + 1 , \begin{aligned} &\hat{\mathbf{z}}^{l}=\mathrm{W-MSA}\left(\mathrm{LN}\left(\mathbf{z}^{l-1}\right)\right)+\mathbf{z}^{l-1}, \\ &\mathbf{z}^{l}=\mathbf{MLP}\left(\mathbf{LN}\left(\mathbf{\hat{z}}^{l}\right)\right)+\mathbf{\hat{z}}^{l}, \\ &\hat{\mathbf{z}}^{l+1}=\mathrm{SW-MSA}\left(\mathrm{LN}\left(\mathbf{z}^{l}\right)\right)+\mathbf{z}^{l}, \\ &\mathbf{z}^{l+1}=\mathbf{MLP}\left(\mathbf{LN}\left(\mathbf{\hat{z}}^{l+1}\right)\right)+\mathbf{\hat{z}}^{l+1}, \end{aligned} z^l=W−MSA(LN(zl−1))+zl−1,zl=MLP(LN(z^l))+z^l,z^l+1=SW−MSA(LN(zl))+zl,zl+1=MLP(LN(z^l+1))+z^l+1,
CNN金字塔中第一层的输出将被输入到1 × 1的卷积中，以生成长度为D '的 ( H / 4 × W / 4 ) (H/4\times W/4) (H/4×W/4)补丁(令牌)。这些补丁经过第一个Swin Transformer块，生成第一个基于注意力的特征映射。跳过连接将先前的激活添加到获得的特征映射中，从而产生最大的分支特征映射p1。接下来，应用补丁合并，将 2 × 2 2\times 2 2×2组相邻的补丁连接起来，应用线性层，在降低分辨率的同时将嵌入维数从D '增加到2D '。类似地，CNN和基于注意力的特征图的高级特征图被融合并馈送到Swin Transformer块中以生成高级输出。后者用最小层特征映射P s表示。
3.1.3 双级融合模块(DLF)
主要的挑战是有效地融合CNN和Swin Transformer的电平特征，同时保持特征的一致性。一种直接的方法是通过解码器直接将CNN的电平和它们匹配的Swin Transformer电平相加，从而获得分割图。然而，这种方法不能保证它们之间的特征一致性，导致性能欠佳。
因此，我们提出了一种新的双水平融合(DLF)模块，它将最小(P s)和最大(P l)水平作为输入，并采用交叉注意机制跨尺度融合信息。
一般来说，较浅的层次有更好的定位信息，当我们接近更深的层次时，语义信息变得更加普遍，更适合解码器部分。面对中间层特征映射对模型精度影响不明显、计算成本大的困境，我们没有考虑在特征融合中使用中间层来节省计算成本。
因此，我们鼓励通过在保留定位信息的同时对最浅( P s P_s Ps)和最后( P l P_l Pl)层进行多尺度表示。
在本文提出的DLF模块中，类令牌的作用非常重要，因为它概括了输入特征的所有信息。我们为每个级别分配一个类令牌，该类令牌来自于该级别规范之上的全局平均池化(GAP)。我们获得类令牌，如下所示:
C L S s = G A P ( N o r m ( P s ) ) C L S l = G A P ( N o r m ( P l ) ) \begin{aligned}CLS^s&=GAP(Norm(P^s))\\CLS^l&=GAP(Norm(P^l))\end{aligned} CLSsCLSl=GAP(Norm(Ps))=GAP(Norm(Pl))
其中 C L S s ∈ R 4 D ′ × 1 CLS^s\in R^{4D^{\prime}\times1} CLSs∈R4D′×1, C L S l ∈ R D ′ × 1 CLS^{l}\in R^{D^{\prime}\times1} CLSl∈RD′×1。然后，在将类令牌传递到转换器编码器之前，将其与相关的级别嵌入连接起来。小电平接S，大电平接LTrasnformer用于计算全局自关注的编码器。值得注意的是，我们还为两个级别的每个标记添加了一个可学习的位置嵌入，然后将它们提供给Trasnformer编码器以学习位置信息
在通过Trasnformer编码器的嵌入后，使用交叉注意模块融合每个级别的特征。具体来说，在融合之前，两个级别的类标记被交换，这意味着一个级别的类标记与另一个级别的标记连接。
然后，每个新的嵌入分别通过该模块进行融合，最后返回到自己的水平。这种与其他关卡令牌的交互使类令牌能够与它们的跨关卡共享丰富的信息。
特别是，图1b显示了小水平的位移。f s(.)首先将CLS投影到p1的维数上，输出记为CLS。CLS 's与p1连接作为键和值，并独立执行计算注意力的查询。由于我们只查询类令牌，所以交叉注意机制在线性时间内运行。最终输出Z s可以用数学方法写成如下:
y s = f s ( C L S s ) + M C A ( L N ( [ f s ( C L S s ) ∥ P l ] ) ) Z s = [ P s ∥ g s ( y s ) ] \begin{aligned} &y^{s}=f^{s}(CLS^{s})+MCA(LN([f^{s}(CLS^{s})\parallel P^{l}])) \\ &Z^{s}=[P^{s}\parallel g^{s}(y^{s})] \end{aligned} ys=fs(CLSs)+MCA(LN([fs(CLSs)∥Pl]))Zs=[Ps∥gs(ys)]
3.2 解码器
在语义FPN[31]的激励下，我们设计了一个解码器，该解码器将ps和p1级的特征结合为一个统一的掩码特征。首先，从DLF模块接收低分辨率和高分辨率特征图P s和P l。ps (H/16, W/16)之后是一个ConvUp块，该块应用3 × 3 Conv、2×双线性上采样、Group Norm[50]和ReLU两个阶段来获得(H/ 4, W/4)分辨率。p1 (H/4, W/4)之后还有一个Conv块，它采用3×3 Conv、群范数和ReLU，并保持在(H/4, W/4)分辨率。处理后的P s和P l的和经过另一个ConvUp块，得到最终统一的H × W特征图。将获得的特征映射通过分割头中的3×3 Conv传递后，生成最终的分割映射。
