HiFormer Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation

[WACV2023] HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation

摘要

由于卷积神经网络的卷积运算的特性，它们在建模长程相关性和空间相关性时受到限制。虽然Transformer最初是为了解决这个问题而开发的，但它们无法捕获低级别的特征。相比之下，局部和全局特征对于密集预测(例如在具有挑战性的环境中进行分割)都是至关重要的。
在本文中，我们提出了一种新的方法HiFormer，该方法有效地连接了CNN和用于医学图像分割的Transformer。具体来说，我们使用开创性的Swin Transformer模块和基于CNN的编码器设计了两个多尺度特征表示。为了确保从上述两种表示中获得的全局和局部特征的良好融合，我们在编码器-解码器结构的跳跃连接中提出了双级融合(Double-Level fusion, DLF)模块。
在各种医学图像分割数据集上的大量实验表明，HiFormer在计算复杂度、定量和定性结果方面优于其他基于CNN、基于Transformer和混合方法。

1 引言

医学图像分割是计算机视觉的主要挑战之一，它提供了详细分析所需的解剖区域的有价值的信息。这些信息可以极大地帮助医生描述损伤，监测疾病进展，并评估是否需要适当的治疗。由于医学图像分析的使用越来越多，高精度和鲁棒分割变得越来越重要。

卷积神经网络(CNN)具有提取图像特征的强大能力，已广泛应用于不同的图像分割任务。随着基于编码器-解码器的网络的兴起，如全卷积网络(FCNs) ，U型结构，例如U-Net及其变体，CNN在医学图像分割任务中取得了显著的成功。在这两种结构中，跳跃连接分别用于体现编码器和解码器路径提供的高级和细粒度特征。尽管CNN模型在各种计算机视觉任务中取得了成功，但由于其接受域有限和固有的归纳偏差，其性能受到限制。上述原因阻止了CNN在图像中构建全局上下文和长期依赖关系，因此限制了它们在图像分割中的性能。

最近，由于Transformer在自然语言处理(NLP)中取得了突出的成功，Vision Transformer已经被开发出来，以缓解CNN在图像识别任务中的缺点。Transformer主要利用一个多头自注意(MSA)机制，该机制可以有效地在token序列和捕获全局上下文之间构建长期依赖关系。Vision Transformer表现出与基于CNN的方法相当的性能，但需要大量数据来泛化，并且具有二次复杂度。有几种方法被提出来解决这些限制。DeiT提出了一种高效的知识蒸馏训练方案，克服了Vision Transformer需要大量数据学习的困难。Swin Transformer和pyramid vision Transformer分别尝试利用基于窗口和空间的约简注意来降低视觉变压器的计算复杂度。

此外，多尺度特征表示最近在Vision Transformer中表现出强大的性能。CrossViT提出了一种新颖的双分支Transformer架构，可提取多尺度上下文信息，并为图像分类提供更细粒度的特征表示。同样，DS-TransUNet提出了一个双分支Swin Transformer来捕获编码器中不同的语义尺度信息，用于医学图像分割任务。HRViT连接多分支高分辨率架构与Vision Transformer进行语义分割。因此，这样的结构可以有效地帮助增强token之间远程关系的建模，并获得更详细的信息。

尽管Vision Transformer能够对全局上下文表示进行建模，但自注意机制会导致缺失低级特征。为了缓解上述问题，已经提出了混合CNN - Transformer方法，利用CNN的局域性和Transformer的长域依赖特性来编码全局和局部特征，特别是在医学图像分割中TransUnet和LeVit - Unet。然而，这些方法都存在一些障碍，阻碍了它们获得更高的性能:1)它们不能在保持特征一致性的情况下有效地结合低级特征和高级特征，2)它们没有正确地使用分层编码器产生的多尺度信息。

在本文中，我们提出了一种新的基于编码器-解码器CNN—Transformer的框架，该框架有效地利用了Transformer的全局长范围关系和CNN的局部特征表示，用于精确的医学图像分割任务。编码器包括三个模块:两个分层CNN和Swin Transformer模块和DLF模块。Swin Transformer和CNN模块都包含三个层次。首先，将输入图像输入CNN模块以学习其局部语义表示。为了弥补全局表示的不足，Swin Transformer模块应用于CNN的浅层特性之上，以捕获长期依赖关系。接下来，利用具有不同窗口大小的Swin Transformer模块的金字塔来学习多尺度交互。为了鼓励特性的可重用性并提供本地化信息，设计了一个跳过连接模块来将CNN的本地特性传输到Transformer块中。然后将最小和最大金字塔级别的结果表示形式输入DLF模块。新提出的DLF模块是一个多尺度Vision Transformer，它使用交叉注意机制融合两个获得的特征图。最后，两个重新校准的特征映射被传递到解码器块以产生最终的分割掩码。我们提出的HiFormer不仅缓解了上述问题，而且在不同的评估指标上也超越了所有的同类产品。我们的主要贡献:

• 一种新的混合方法，融合了Transformer的远程上下文交互和CNN的本地语义信息。

• DLF模块在粗粒度和细粒度特征表示之间建立有效的特征融合。

• 实验结果证明了HiFormer算法在医学图像分割数据集上的有效性和优越性。

2 相关工作

2.1 基于CNN的分割网络

卷积神经网络被认为是不同计算机视觉任务的事实上的标准。CNN取得优异成绩的一个领域是图像分割，其中为每个像素分配类别标签。Long等表明，完全卷积网络(FCNs)可以在没有完全连接层的情况下分割图像。考虑到常规FCNs其中卷积层按顺序堆叠)的输出通常比较粗糙，因此提出了融合不同层输出的其他模型。为了改善FCN有限的感受野，已经引入了几种方法，包括扩张卷积和上下文建模。CNN模型在医学成像任务中表现出色。在U-net问世后，其他研究者开始关注U型编解码器结构的应用。[46]中使用U-net扩充过完备网络，Unet++中通过在模块之间增加密集跳过连接，重新设计了编解码器架构。这种结构在不同的医学领域得到了进一步的改进和应用。

2.2 Vision Transformer

继NLP中Transformer的巨大成功之后，Dosovitskiy等提出了Vision Transformer (ViT)，它通过使用自注意机制学习全局信息，在图像分类任务中实现了最先进的性能。为了提高Vision Transformer的效率，减少对大型数据集的依赖以实现泛化，人们引入了几种Vision Transformer的衍生物。

此外，还提出了许多方法，重点是多尺度表示，通过从不同尺度提取信息来提高精度和效率。灵感来自CNN的金字塔结构，PVT是最早推出的金字塔Vision Transformer。随后，Swin Transformer提出了一种分层Vision Transformer，使用一种有效的移位窗口方法在局部计算自注意。CrossViT建议使用双分支Vision Transformer，然后使用交叉注意模块，以便在线性时间内执行更丰富的特征表示。Vision Transformer在其他视觉任务中也显示了令人印象深刻的结果，包括[58,20]，它为目标检测提供了基于Transformer的端到端模型，以及[43,24]用于语义和实例分割

2.3 用于医学图像分割的Transformer

尽管CNN模型的结果令人鼓舞，但由于其接受域有限，这种方法通常表现出建模长期依赖关系的限制，从而产生较弱的性能。最近，基于Transformer的模型在医学图像分割中比CNN模型得到了显著的普及。Swin- Unet和DS-TransUNet提出了基于Swin Transformer的U型架构的纯Transformer模型，用于二维分割。除了完整的Transformer模型外，TransUNet还利用CNN和Transformer来捕获低级和高级特征。UNETR使用基于Transformer的编码器嵌入输入的3D补丁，并使用基于CNN的解码器实现最终的3D分割结果。以往的研究大多利用CNN进行特征提取，缺乏全局特征，或者利用Transformer进行特征提取，局部特征表示有限;这将导致不包含丰富信息的无效特征图。在混合工作中，简单的特征融合机制不能保证特征在不同尺度之间的一致性。在多尺度表示的激励下，我们提出了一种基于CNN -Transformer的架构HiFormer，该架构有效地融合了全局和局部信息，并利用一种新的基于Transformer的融合方案来保持二维医学图像分割任务的特征丰富性和一致性。

3 方法

本节将概述所提议的HiFormer。如图1a所示，我们提出的体系结构提供了一个端到端训练策略，该策略集成了来自Swin Transformer的全局上下文表示和来自编码器中CNN模块的局部代表特征。然后使用双层融合模块(DLF)获得更丰富的特征表示。之后，解码器输出最终的分割映射。

图1:(a) HiFormer的概述。HiFormer由一个分层的CNN-Transformer特征提取器模块组成;第一层和最后一层的输出通过DLF特征融合模块提供。随后，解码器使用DLF的输出生成精确的分割映射。在图中，蓝色和橙色块分别表示Swin Transformer和CNN级别。(b)交叉注意概述。小级别的类标记CLS’^s首先被投影用于维度对齐，然后追加到p^l。结果嵌入作为键和值执行。此外，还利用CLS进行查询。最后计算注意力和反投影，得到Z^s。这个过程也可以扩展到更大的层面。

3.1 编码器

如图1a所示，所提出的编码器由CNN和Swin Transformer两个层次模型组成，DLF模块丰富检索到的特征，并准备将其输入解码器。由于单独使用CNN或Transformer会导致局部或全局特征被忽略，从而影响模型的性能，因此我们首先利用CNN的局域特征来获得局部特征。在这里，CNN和Swin Transformer分别包含三个不同的级别。我们通过跳跃连接将每个层次的局部特征转移到相应的Swin Transformer的层次，以获得通用表示。然后将每个传输的CNN级别与其并行Transformer级别相加，并通过Patch merge模块生成分层表示(见图1a)。我们利用分层设计来利用多尺度表示。最大和最小的级别进入DLF模块，以交换来自不同规模的信息并生成更强大的功能。

在下面，我们将深入详细地讨论我们的CNN、Swin Transformer和DLF模块。

3.1.1 CNN模块

所提出的编码器首先使用CNN作为特征提取器来构建不同分辨率的中间CNN特征图金字塔。取空间维度为H和W，通道为C的输入图像X∈R^H×W×C，首先送入CNN模块。CNN模块由三个层次组成，其中使用Conv 1 × 1将跳过连接连接到相关Transformer的层次，以补偿Transformer的低层次缺失信息并恢复局部空间信息。

3.1.2 Swin Transformer模块

普通的Transformer编码器块由两个主要模块组成:一个多头自注意(MSA)和一个多层感知器(MLP)。Vanilla Transformer由N个相同的Transformer编码器块组成。在每个块中，在MSA和MLP块之前，应用LayerNorm (LN)。此外，还通过跳过连接将激活的副本添加到MSA或MLP块的输出中。使用标准MSA的普通ViT的一个主要问题是它的二次复杂度，这使得它在执行图像分割等高分辨率计算机视觉任务时效率低下。为了克服这一限制，Swin Transformer引入了WMSA和SW-MSA。

Swin Transformer模块包括两个连续修改的Transformer块;将MSA块替换为基于窗口的多头自注意(W-MSA)和移位的基于窗口的多头自注意(SW-MSA)。在W-MSA模块中，自注意作用于大小为M × M的局部窗口，WMSA模块具有线性复杂度;然而，由于没有跨窗口的连接，它的建模能力有限。为了缓解这一问题，引入了SW-MSA，它利用了与W-MSA模块的输入相比移位的窗口配置;这是为了确保我们有跨窗口连接。此过程描述在下面公式：

CNN金字塔中第一级的输出将被输入到1 × 1卷积中，以生成长度为D’的(H/4 × W/4)块(标记)。这些Patch通过第一个Swin Transformer块，生成第一个基于注意力的特征映射。跳过连接将先前的激活添加到获得的特征映射中，从而得到最大的分支特征映射P^l。接下来，应用patch合并层，将2 × 2组相邻的patch进行拼接，应用线性层，在降低分辨率的同时将嵌入维数从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^l)层进行多尺度表示。

在提议的DLF模块中，class token扮演着重要的角色，因为它总结了输入特征的所有信息。我们为每个级别分配一个class token，该class token派生自该级别规范上的全局平均池(GAP)。我们获得class token，如下所示:

其中CLS^s∈R^{4D ’ ×1}, CLS^l∈R^{D ’ ×1}。然后，在传递到Transformer编码器之前，class token与相关的级别嵌入连接。小的块接S Transformer，大的块接L Transformer用于计算全局自我注意的编码器。值得注意的是，我们还为两个级别的每个标记添加了一个可学习的位置嵌入，然后将它们交给Transformer编码器以学习位置信息。

在通过Transformer编码器传递嵌入后，使用交叉注意模块融合每个级别的特征。具体来说，在融合之前，两个级别的class token被交换，这意味着一个级别的class token与另一个级别的class token连接。然后，每个新的嵌入分别通过模块进行融合，最后反投影到它自己的水平。这种与其他级别class token的交互使class token能够与其跨级别共享丰富的信息。

特别是小等级的位移如图1b所示。f^s(.)首先将CLS^s投影到P l的维数上，输出记为CLS^'s。CLS^'s与P^l的连接作为键和值，并独立执行计算注意力的查询。由于我们只查询class token，交叉注意机制在线性时间内运行。最终的输出Z ^s可以用数学形式写成:

3.2 解码器

在语义FPN的激励下，我们设计了一个将P^s和P^l层的特征组合成统一掩码特征的解码器。首先，从DLF模块接收低分辨率和高分辨率的特征图P^s和P^l。P^s (H/16, W/16)之后是一个ConvUp块，该块应用3 × 3 Conv、2×双线性上采样、Group Norm和ReLU两个阶段来获得(H/ 4, W/4)分辨率。P l (H/4, W/4)后面还跟着一个Conv块，它采用3×3 Conv，Group Norm和ReLU，并保持在(H/4, W/4)分辨率。经过处理的P^s和P^l的总和经过另一个ConvUp块，以获得最终统一的H × W特征图。将获取的特征图通过分割头中的3×3 Conv后，生成最终的分割图。

4 实验

4.1 数据集

**Synapse Multi-Organ Segmentation:**首先，我们评估HiFormer在基准Synapse多器官分割数据集上的性能。该数据集包括30例病例3779张轴向腹部临床CT图像，每个CT体积涉及85 ~ 198个512 × 512像素的切片，体素空间分辨率为([0.54 ~ 0.54]× [0.98 ~ 0.98] × [2.5 ~ 5.0]) mm³。

**Skin Lesion Segmentation: **我们在皮肤病变分割数据集上进行了广泛的实验。具体来说，我们使用ISIC 2017数据集，其中包括2000张皮肤镜图像用于训练，150张用于验证，600张用于测试。此外，我们采用ISIC 2018，并根据文献将数据集划分为训练集、验证集和测试集。此外，使用PH2数据集，这是一个皮肤镜图像数据库，用于分割和分类任务。

**Multiple Mylomia Segmentation:**我们还评估了SegPC 2021提供的多发性骨髓瘤细胞分割的方法。挑战数据集包括一个包含290个样本的训练集和一个包含200个样本和277个样本的验证和测试集。

4.2 实现细节

我们在PyTorch中实现了我们的框架，并在一个具有24 GB内存的Nvidia RTX 3090 GPU上进行了训练。输入图像大小为224 × 224，训练时我们将批大小和学习率分别设置为10和0.01。此外，我们使用ImageNet上预先训练好的CNN和Swin Transformer模块的权重来初始化它们的参数。我们的模型使用SGD优化器进行优化，动量为0.9，权重衰减为0.0001。此外，在训练过程中使用翻转和旋转等数据增强来提高多样性。

表1描述了建议模型的最终配置。

表1:建议的模型配置。WS表示窗口大小，D '表示嵌入尺寸，r表示Transformer块使用的MLP扩展比。DLF模块中的头的数量对于这两个级别是相同的。

4.3 评价结果

我们在每个实验的评估指标方面采用了特定于任务的范例。具体来说，这些指标包括Dice评分、95%豪斯多夫距离(HD)、敏感性和特异性、准确性和mIOU。为了确保无偏见的比较，我们将HiFormer与CNN和基于Transformer的方法进行对比，以及在两者合并的基础上制定的模型。

4.3.1 腹部多器官分割结果

在8个腹部器官的平均骰子相似系数(DSC)和平均豪斯多夫距离(HD)方面，该建议与以前的最先进(SOTA)方法的比较如表2所示。HiFormer的性能优于基于CNN的SOTA方法。与其他基于Transformer的模型相比，我们的HiFormer-B在两个评估指标上都表现出了优越的学习能力，与TransUnet和Swin-Unet相比，我们的Dice得分分别提高了2.91%和1.26%，平均HD分别下降了16.99和6.85。具体而言，HiFormer稳居大部分器官的分割的前列，特别是胃、肾和肝的分割。可以观察到，HiFormer在平均HD方面比其他方法有明显的优势。此外，参数数量方面的效率如表2所示，这将在接下来的章节中讨论。图2给出了结果的一个定性的典型例子。实验结果表明，该方法能够准确分割精细复杂结构，并输出更精确的分割结果，对复杂背景具有更强的鲁棒性。

表2:该方法在Synapse数据集上的比较结果。蓝色表示最好的结果，红色表示次好的结果

图2:该方法在Synapse数据集上的分割结果。红色矩形标识器官区域，我们提出的方法的优越性可以清楚地看到。

4.3.2 Skin Lesion Segmentation的结果

ISIC 2017、ISIC 2018和PH2皮肤病变分割任务基准与领先方法的比较结果如表3所示。我们的HiFormer在大多数评估指标上比其他竞争对手表现得更好。具体而言，HiFormer在不同数据集上的优势凸显了其令人满意的泛化能力。我们还在图3中展示了皮肤病变分割结果的视觉比较，这表明我们提出的方法能够捕获更精细的结构并生成更精确的轮廓。具体而言，如图3所示，我们的方法在边界区域的性能优于TMU-Net等混合方法。此外，如图3所示，与纯基于Transformer的方法(如SwinUnet)相比，HiFormer对噪声项具有鲁棒性，后者由于缺乏局域建模而性能下降。这种卓越的性能是通过将Transformer和CNN用于建模全局关系和局部表示而实现的。

表3:所提方法与SOTA方法在皮肤病变分割基准上的性能比较。蓝色表示最好的结果，红色表示次好的结果。

图3:不同方法在ISIC2017皮肤病变分割数据集上的视觉比较。真实边界以绿色显示，预测边界以蓝色显示。

4.3.3 多发性骨髓瘤分割结果

在表4中，我们包括了基于平均IoU度量的结果。在我们测试的所有配置中，HiFormer结构始终优于挑战排行榜。

另外，本文HiFormer的部分分割输出如图4所示。如图所示，我们的预测很好地适应了所提供的GT面具。HiFormer的主要优势之一是它能够建模多尺度表示。它抑制了背景噪声，这是高度重叠背景的数据集(如SegPC)的情况。简单地说，HiFormer超越了只有局部信息建模能力的基于cnn的方法和基于变压器的方法，这些方法在边界区域表现不佳。

4.3.4 模型参数比较

在5中，我们将所提方法的参数数量与医学图像分割模型的参数数量进行了比较。与其他作品相比，我们的轻量级HiFormer在模型复杂性方面表现出了极大的优势，同时获得了出色的或相当的性能。

5 消融实验

不同CNN主干的比较：

我们首先研究不同的CNN主干的贡献。具体来说，我们使用ResNet和DenseNet的变体作为卷积架构的两个先验艺术。如表7所示，使用ResNet主干可以获得最佳性能。此外，我们已经看到更大的CNN主干并不一定会带来性能提升(参见表7中的第3行和第4行)，这让我们了解到使用ResNet50架构作为默认架构

DLF模块的影响:

接下来，我们评估DLF模块对分割性能的重要性。

表6的实验结果揭示了DLF模块在编码和解码过程中不可忽视的作用。具体来说，DLF模块分别为骰子分数和HD带来了显著的改进(3.24%和2.18%)。通过交叉注意机制，DLF模块帮助网络融合全局和局部特征。结果表明，将CNN与Transformer相结合有助于病灶的分割。

不同DLF模块配置上的消融:表8显示了不同DLF模块配置的性能。我们测试正面数目的不同值的大等级和小等级的MLP深度(S和L)以及Transformer模块MLP块中的MLP膨胀比®。我们观察到(S, L)的(2,1)对和两个电平的6个头效果最好。如A行所示，增加正面的数量并不一定能提高性能。此外，膨胀比®对性能也有重要影响。与C行相比，r翻倍导致DSC增加1.04%，HD下降1.82%。

特征一致性的消融：我们进行了两个实验来测量和揭示特征一致性，并在SM中详细讨论了它们。首先，我们展示了涉及DLF模块前后各层的特征可视化(SM，图1-2)。第二个实验证明了应用每个模块如何有助于特征一致性(SM，表3)。总的来说，每个模块在提供更一致的特征方面的贡献可以从结果中推断出来

6 讨论

我们在不同医学图像分割数据集上的综合实验证明了我们提出的HiFormer模型与CNN和基于Transformer的方法相比的有效性。我们的方法的主要进步有两个方面。其设计的第一个合理性是在网络的浅层将CNN和Transformer结合在一起。其次，跳过连接模块提供了特性可重用性，并将CNN本地特性与Transformer模块提供的全局特性混合在一起。定量看待HiFormer网络的五大挑战数据集表明，它可以很好地进行分割，在大多数情况下超过了SOTA方法。从视觉分析的角度来看，图2展示了对肝脏和肾脏等器官的无噪声分割，这也与定量基准一致。相比之下，我们的模型在某些情况下(例如主动脉)获得了失效案例，这与数值结果再次一致。此外，可以看出低对比度的皮肤图像仍然给我们的模型带来了很大的困难。总的来说，HiFormer已经显示出有效学习医学图像中所代表的关键解剖关系的潜力。在模型参数方面，与其他复杂模型相比，HiFormer模型是一个轻量级模型，在医学图像分割中存在严重的问题。

7 结论

噪声分割，这也与定量基准一致。相比之下，我们的模型在某些情况下(例如主动脉)获得了失效案例，这与数值结果再次一致。此外，可以看出低对比度的皮肤图像仍然给我们的模型带来了很大的困难。总的来说，HiFormer已经显示出有效学习医学图像中所代表的关键解剖关系的潜力。在模型参数方面，与其他复杂模型相比，HiFormer模型是一个轻量级模型，在医学图像分割中存在严重的问题。

8 结论

在本文中，我们介绍了一种新的基于混合CNN -Transformer的医学图像分割方法HiFormer。具体来说，我们将从Swin Transformer模块获得的全局特征与基于CNN的编码器的局部表示结合起来。然后，使用DLF模块，我们可以更好地融合来自上述表示的特征。我们实现了优于基于CNN、基于Transformer和混合模型的性能，这表明我们的方法在保持低级特征的细节和建模长期交互方面取得了平衡。