MICCAI 2022 | PHTrans: Parallelly Aggregating Global and Local Representations for Medical Image Seg

MICCAI 2022 | PHTrans: Parallelly Aggregating Global and Local Representations for Medical Image Segmentation

MICCAI 2022 并行聚合医学图像分割的全局和局部表示

0 Abstract

Transformer在计算机视觉领域的成功引起了医学影像界越来越多的关注。特别是在医学图像分割方面，已经出现了许多基于卷积神经网络(CNNs)和Transformer的混合架构，并取得了令人瞩目的性能。然而，大多数将模块化Transformer嵌入cnn的方法都难以发挥其全部潜力。在本文中，我们提出了一种新的混合架构PHTrans用于医学图像分割，该架构将Transformer和CNN在主要构件中并行混合，从全局和局部特征中生成分层表示并自适应聚合，旨在充分利用它们的优点以获得更好的分割性能。具体来说，PHTrans遵循U形编码器-解码器设计，并在深层阶段引入并行混合模块，其中卷积块和改进的3D Swin Transformer分别学习局部特征和全局依赖关系，然后通过**sequence-to-volume (序列到体积)**的操作统一输出的维度以实现特征聚合。在多器官分割（MALBCV）和自动心脏诊断挑战（ACDC）数据集上的大量实验结果证实了其有效性，始终优于SOTA方法。

1 Introduction

医学图像分割旨在提取和量化生物组织/器官图像中感兴趣的区域，这些区域对于疾病诊断、术前规划和干预至关重要。得益于深度学习优异的表示学习能力，卷积神经网络在医学图像分析中取得了巨大的成功。出现了许多优秀的网络模型(如U-Net、3D U-Net、Attention U-Net)，不断刷新各种分割任务的性能上限。尽管获得了极具竞争力的结果，但基于CNN的方法由于固有的**归纳偏差(如局部性和平移不变性)**而缺乏建模长期依赖关系的能力。一些研究人员通过增加卷积核的大小、使用空洞卷积（Yu, F., Koltun, V.: Multi-scale context aggregation by dilated convolutions. arXiv preprint arXiv:1511.07122 (2015)）和嵌入自注意机制来缓解这个问题。然而，只要卷积运算仍然是网络架构的核心，就不能从根本上解决缺乏全局信息的问题。

Transformer，完全依靠注意机制来建模全局依赖关系，而不需要任何卷积操作，已经成为一种替代架构，在大规模数据集上进行预训练的情况下，在计算机视觉(CV)中提供了比CNN更好的性能。其中，Vision Transformer (ViT)将图像分割为一系列tokens，并使用堆叠的Transformer块对其全局关系进行建模，这对CV领域产生了革命性的影响。Swin Transformer可以在可移动窗口中生成计算复杂度较低的分层特征表示，在各种CV任务中实现最先进的性能。然而，医学图像数据集的规模远远小于上述工作中预训练的数据集(例如ImageNet-21k和JFT-300M)，因为医学图像并不总是可用的，需要专业的标注。因此，Transformer在医学图像分割方面的性能并不理想。同时，许多由CNN和Transformer结合而来的混合结构已经出现，它们提供将相互的优点融合已逐渐成为医学图像分割的折中解决方案，无需在大数据集上进行预训练。

本文总结了医学图像分割中几种流行的基于Transformer和CNN的混合架构。这些混合体系结构将Transformer添加到以CNN为骨干的模型中，或者替换体系结构的部分组件。例如：

1. UNETR和Swin UNETR采用了编码器-解码器结构，其中编码器由自注意和多层感知器(即Transformer)构建的层叠块组成，而解码器则是堆叠的卷积层，如图1(a)所示。
2. TransBTS和Trans UNet在编码器和解码器之间引入了一个由CNN组成的Transformer，如图1(b)所示。
3. MISSFormer和CoTr通过Transformer桥接了从编码器到解码器的所有阶段，而不是仅连接相邻阶段，这捕获了多尺度全局依赖关系，见图1©。
4. nnFormer将Transformer和卷积块交织成混合模型，其中卷积编码精确的空间信息，self-attention捕获全局上下文，如图1(d)所示。

从图1可以看出，这些体系结构从宏观角度实现了Transformer和CNN的串行组合。然而，在串行组合中，卷积和自注意不能贯穿整个网络架构，难以连续地对局部和全局表示进行建模，因此不能充分利用它们的潜力。

在本文中，我们提出了一种用于医学图像分割的并行混合Transformer (PHTrans)，其中主要构建块由CNN和Swin Transformer组成，以同时聚合全局和局部表示，见图1(e)。在PHTrans中，我们将标准Swin Transformer扩展到3D版本，方法是提取划分体积的3D补丁并构建3D自注意机制。鉴于Swin Transformer的层次性可以方便地利用U-Net等先进的密集预测技术，我们遵循成功的u型架构设计，引入sequence-to-volume (序列到体积)的转换操作，实现Swin Transformer和CNN在一个块内的并行组合。与串行混合架构相比，PHTrans可以独立构建分层的局部和全局表示，并在每个阶段融合它们，充分利用CNN和Transformer的潜力。大量的实验证明了该方法在各种医学图像分割任务中优于其他竞争方法。

2 Method

2.1 总体架构

PHTrans结构的概述如图2(a)所示。PHTrans采用U型编码器和解码器设计，主要由纯卷积模块和并行混合模块组成。我们的初衷是构造一个由Transformer和CNN组成的完全混合的体系结构，但是由于自注意机制的高计算复杂度，Transformer不能直接接收以像素作为tokens的输入。在我们的实现中，引入了级联卷积块和下采样操作来减小空间大小，从而逐步提取低层次的具有高分辨率的特征，获得精细的空间信息。同样，这些纯卷积模块也部署在解码器中，通过上采样恢复原始图像的维数

给定一个输入体积x∈R^{H×W ×D}，其中H、W、D分别表示高度、宽度、深度，我们首先利用几个纯卷积模块得到特征映射
$$f\in R^{\frac{H}{2^{N}1}\times \frac{W}{2^{N}1}\times \frac{D}{2^{N}1}\times 2^{N}1C}$$
其中N1和C分别表示模块数和基通道数。然后，利用Transformer和CNN组成的并行混合模块对局部特征和全局特征的层次表示进行建模。以
$$\frac{H}{2^{N1+N2}}\times \frac{W}{2^{N1+N2}}\times \frac{D}{2^{N1+N2}}$$
为输出分辨率，2^N1+N2C为通道数量，重复计算N2次。与编码器相对应的对称解码器同样基于纯卷积模块和并行混合模块，通过跳跃连接和加法操作融合编码器的语义信息。此外，我们在训练过程中在解码器的每个阶段都使用深度监督，导致总共N₁+N₂个输出，其中应用了由交叉熵和Dice损失组成的联合损失。PHTrans的架构简单多变，每个模块的数量可以根据医学图像分割任务进行调整，即N₁, N₂, M₁和M₂。其中，M₁和M₂为并行混合模块中Swin Transformer块和卷积块的个数。

2.2 并行混合模块

并行混合模块部署在PHTrans的深层，其中Trans&Conv块作为其核心，通过CNN和Swin Transformer实现局部和全局表示的分层聚合。

2.2.1 Trans&Conv块。

将缩小比例的特征图分别送入Swin Transformer (ST)块和convolution (Conv)块。我们在ST块的开始和结束分别引入了**volume -to- sequence (V2S)和sequence -to- volume (S2V)**操作来实现体积和序列的转换，使其与Conv块产生的输出的维空间一致。具体来说，V2S用于将整个体积(3D图像)重塑为具有窗口大小的3D补丁序列。S2V是相反的运算。

如图2(b)所示，一个ST块包括一个基于移位窗的多头自我注意（MSA）模块，然后是一个2层的
MLP，中间有一个GELU激活函数。在每个MSA模块和每个MLP模块之前加LN (LayerNorm)层，在每个模块之后加残差连接。在M1个连续的ST块中，将具有常规（Regular）和移位（Shift）窗口配置的MSA，即W-MSA和SW-MSA，交替嵌入到ST块中，实现跨窗口连接，同时保持不重叠窗口的高效计算。

对于医学图像分割，我们将标准ST块修改为3D版本，在局部3D窗口内计算自注意，这些窗口以不重叠的方式均匀地划分体积。假设x∈R^H×W×S×C是ST块的输入，首先将其reshape为N×L×C，其中N和L = Wh × Ww × Ws分别表示三维窗口的个数和维数。每个头的自注意计算公式为:
$$Attention(Q,K,V)=SoftMax(\frac{Q{K}^T}{\sqrt{d}}+B)V,$$
其中Q、K、V∈R^L×d为query、key和value矩阵，d为query/key维数，B∈R^L×L为相对位置偏差。我们参数化了一个较小的偏置矩阵B∈R^{(2Wh−1)×(2Ww−1)×(2Ws−1)}，B中的值取自B帽。

卷积块以3 × 3 × 3卷积层、GELU非线性和实例归一层为单位重复M2次。卷积块的配置简单而灵活，任何现成的卷积网络都可以应用。最后，我们通过加法运算将ST块和Conv块的输出融合。编码器中Trans&Conv块的计算过程可以概括为：
$$y_i=S2V(S{T}^{M1}(V2S(x_{i-1})))+Con{v}^{M2}(x_{i-1})$$
其中xi−1是编码器第i−1级的下采样结果。在解码器中，除了跳过连接之外，我们还使用加法操作补充来自编码器的上下文信息。因此，解码器中的Trans&Conv块可以表述为:
$$z_i=S2V(S{T}^{M1}(V2S(x_{i+1}+y_i)))+Con{v}^{M2}([x_{i-1},y_i])$$
其中xi+1是解码器第i+1级的上采样结果，yi是编码器第i级的输出。

2.2.2 下采样和上采样。

下采样包含一个有步长的卷积操作和一个实例归一化层，其中通道数减半，空间大小增加一倍。类似地，上采样是一个有步长的反卷积层，然后是一个实例归一化层，它使特征映射通道的数量增加了一倍，空间大小减半。在所有维度中，步幅通常设置为2。然而，当三维医学图像是各向异性时，相对于特定维度的stride设置为1。

3 实验

3.1 数据集

The Multi-Atlas Labeling Beyond the Cranial Vault (BCV) 多器官分割任务 包括30例病例，3779张轴向腹部临床CT图像。与nnformer类似，数据集被分为18个训练样本和12个测试样本。并以平均dice相似系数(DSC)和平均Hausdorff Distance (HD)作为评价指标，对我们的方法在8个腹部器官(主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、脾脏、胃)上进行评价。

自动心脏诊断挑战(ACDC)数据集是使用MRI扫描仪从不同的患者中收集的。对于每个患者的MR图像，左心室(LV)，右心室(RV)和心肌(MYO)被标记。按照nnformer，70个样本分成训练集，10个样本分成验证集，20个样本分成测试集。平均DSC被用来评估我们在这个数据集上的方法。

3.2 实现细节

为了公平比较，我们使用nnUNet的代码框架来评估PHTrans与CoTr和nnFormer相同的性能。

所有实验均在nnUNet的默认配置下进行。在PHTrans中，我们经验地将超参数[N1,N2,M1,M2]设置为[2,4,2,2]，并采用nnU-Net的跨步策略进行下采样和上采样，无需精心设计。通道C的基数为24，不同编码器阶段使用的多头自注意的头数为[3,6,12,24]。对于BCV数据集和ACDC数据集，我们将ST块中3D窗口[Wh,Ww,Ws]的大小分别设置为[3,6,6]和[2,8,7]。

在训练阶段，我们分别从BCV和ACDC数据集中随机裁剪大小为48×192×192和16×256×224的子卷作为输入。我们在PyTorch 1.9下实现了PHTrans，并使用单个GeForce RTX 3090 GPU进行了实验。

3.3 结果

3.3.1 与SOTA比较

我们将PHTrans的性能与之前最先进的方法进行了比较。除了介绍中提到的混合架构，它还包括，LeViT-U-Net， Swin-Unet和nnU-Net。此外，我们通过将UNETR编码阶段的ViT修改为Swin Transformer来再现Swin UNETR，并以同样的方式评估了UNETR和Swin UNETR的性能。并使用了与我们相同的数据集分区。

在BCV数据集上的分割结果如表1所示。

我们的PHTrans实现了最佳性能88.55% (DSC↑)和8.68 (HD↓)，在平均DSC上超过了之前的最佳模型0.8%，HD超越了1.15。

图3中的代表性样本证明了PHTrans识别器官细节的成功，例如，“胃”在第1、2行，“左肾”在第2行。ACDC数据集上的分割结果如表2所示。

同样，PHTrans与其他最先进的方法相比，具有最高的平均DSC，这是显而易见的。值得一提的是，Swin-Unet、TransUNet、LeViT和nnFormer使用ImageNet上预训练的权重来初始化它们的网络，而PHTrans是在两个数据集上从头开始训练的。

此外，我们还比较了参数的数量和flop的数量，以评估BCV实验中3D方法nnformer、CoTr、nnU-Net、UNETR、Swin UNETR和PHTrans的模型复杂性。如表3所示，PHTrans参数较少(36.3M)，其FLOPs (187.4G)显著低于CoTr、nnU-Net和Swin UNETR。综上所述，PHTrans在BCV和ACDC数据集上的结果充分证明了其出色的医学图像分割和泛化能力，同时保持了适度的模型复杂度。

分割表现如下

3.3.2 消融实验

以改进后的3D Swin-Unet为基准，逐步整合PHTrans的组件，探索不同组件对模型性能的影响。表4提供了消融研究的定量结果。“+PCM”表示使用堆叠的纯卷积模块，而不是跨步卷积操作进行patch分区，而“w/o PCM表示相反。“w/o ST”意味着PHTrans中的并行混合模块删除了Swin Transformer块，从而产生与nnU-Net类似的架构。

从这些结果可以看出，PCM能够提高3D Swin-Unet和PHTrans的性能，这得益于PCM能够在前几个阶段捕获细粒度的细节。此外，与单一架构相比，PHTrans带来了更显著的性能提升，其平均DSC分别比“3D Swin-Unet+PCM”和“PHTrans w/o ST”分别高出3.6%和0.84%，HD分别高出10.84和5.69。结果表明，使用CNN和Transformer并行组合策略聚合全局和局部表示是有效的。

3.3.3 讨论

在PHTrans中，应用了一个普通的Swin Transformer和简单的卷积块，这表明与最先进的Transformer和CNN块相比，显著的性能提升来自于并行混合架构设计。此外，PHTrans不是预训练的，因为到目前为止还没有足够大的3D医学图像的通用数据集。从以上考虑，未来我们将精心设计Transformer和CNN块，探索如何端到端预训练Transformer，进一步提高分割性能。

4 结论

在本文中，我们提出了一种基于Swin Transformer和CNN的并行混合架构(PHTrans)，用于精确的医学图像分割。与其他混合架构将模块化Transformer嵌入到CNN中不同，PHTrans在整个模型中构建了由Swin Transformer和CNN组成的混合模块，不断聚合来自全局和局部特征的层次表示，充分发挥两者的优势。

在BCV和ACDC数据集上的大量实验表明，我们的方法优于几种最先进的替代方法。作为一种通用架构，PHTrans非常灵活，可以被现成的卷积和Transformer模块所取代，这为更多下游医学图像任务开辟了新的可能性。