[深度学习论文笔记]TransBTSV2: Wider Instead of Deeper Transformer for Medical Image Segmentation

TransBTSV2: Wider Instead of Deeper Transformer for Medical Image Segmentation
TransBTSV2:用于医学图像分割的宽Transformer代替深Transformer
Published : Jan 2022
论文：https://arxiv.org/abs/2201.12785
代码：https://github.com/Wenxuan-1119/TransBTS

摘要：
  Transformer得益于利用自注意机制的全局(远程)信息建模，近年来在自然语言处理和计算机视觉领域取得了成功。卷积神经网络能够捕获局部特征，但无法从全局特征空间建立显式的长距离依赖关系模型。然而，局部和全局特征对于密集预测任务至关重要，特别是在三维医学图像分割中。在本文中，作者利用Transformer在3D CNN中进行3D医学图像的体积分割，并提出了一种新的基于编码器-解码器结构的网络TransBTSV2。与原来的TransBTS不同，提出的TransBTSV2不局限于脑肿瘤分割(BTS)，而是专注于一般医学图像的分割，为医学图像的体积分割提供了一个强大而高效的3D基线。TransBTSV2是一种混合的CNN-Transformer架构，可以在不进行任何预训练的情况下实现医学图像的精确分割。通过对Transformer内部结构的重新设计，以及引入的可变形瓶颈模块，实现了具有卓越性能的高效架构。在四个医学图像数据集(BraTS 2019、BraTS 2020、LiTS 2017和KiTS 2019)上的大量实验结果表明，与最先进的脑瘤、肝瘤和肾瘤分割方法相比，TransBTSV2取得了类似或更好的结果。

问题动机：
  癌症作为最常见的临床疾病之一，每年都造成大量的死亡。医学图像的精确测量可以帮助医生做出准确的诊断和进一步的治疗计划。医学图像分割的目的是从背景中识别出肿瘤，勾画出器官的不同子区域，通过为医学图像中的每个像素分配预定义的类别标签，如磁共振成像(MRI)和计算机断层扫描(CT)。传统上，病变区域的划定主要由临床医生依据临床经验进行，费时且容易出错。因此，为了提高临床诊断的准确性和效率，促进医学图像自动分割的发展至关重要。
  与自然图像不同，医学图像的一些独特特性使得自动分割任务非常具有挑战性。首先，不同肿瘤的形状和纹理差异很大，通过直接匹配很难找到它们之间的共性。此外，有些病灶边界不清晰，使得这些肿瘤的轮廓难以准确分离。此外，肿瘤结构在大小、扩展和位置上的差异也阻碍了分割算法使用强先验知识。因此，对医学图像进行自动、准确的分割具有很大的挑战性。
  为了解决上述医学图像分割问题，人们提出了许多方法，包括统计形状模型、基于轮廓的方法和基于机器学习的方法。然而，这些方法严重依赖于手工制作的特征，并且特征表示能力有限。近年来，卷积神经网络(Convolutional Neural Networks, CNNs)在计算机视觉建模中占据了主导地位。由于局部特征是医学图像中清晰分割边界和病灶的关键，因此快速采用cnn进行医学图像分割任务。特别是U-Net采用了带跳连接的对称编解码器结构来细化细节保存。U-Net通过在解码器中进行跳跃连接和连续上采样操作，实现了多级特征的融合，恢复出更精确的分割边缘，成为医学图像分割的主流架构。在接下来的研究中，U-Net++和Res-UNet等UNet变体进一步提高了分割性能。然而，尽管基于cnn的方法具有突出的表示能力，但由于卷积核感受野有限，cnn很难对显式的长距离依赖进行建模。卷积运算的这种内在限制阻碍了cnn学习全局语义信息，而这对于像分割这样的密集预测任务来说是至关重要的。
  现有研究受自然语言处理(NLP)中注意机制的启发，将注意机制与CNN模型融合，克服了无法对远程依赖进行建模的缺陷。注意机制的核心思想是建立特征表示之间的长期关系模型，并基于注意权重更新特征映射。与卷积不同，单个注意层能够从整个特征空间捕获全局语义信息。例如，非局部神经网络的目标是基于自注意机制捕捉特征映射中的长距离依赖关系，但存在较高的内存和计算成本。由于建模显式的远程依赖关系对于医学图像中器官和肿瘤的精确分割至关重要，注意机制也被广泛应用于医学图像分割，取得了显著的效果。
  基于自注意机制，Transformer能够捕获序列中任意位置之间的关系，用于对序列到序列任务中的长期依赖关系进行建模。与注意力机制相比，Transformer在全局上下文建模方面更强大，其独特的架构使其具有更强的特征表示能力，更适合克服医学图像分割的困难。最近，基于transformer的框架已经在各种计算机视觉任务上达到了最先进的性能。例如，Vision Transformer (ViT)架构，通过对大规模数据集的预训练，在图像分类方面取得了出色的性能。之后，许多基于transformer的架构被用于密集的预测任务。
  尽管基于Transformer的方法已经在各种计算机视觉任务中取得了很好的效果，但Transformer的训练和部署仍然是一个巨大的挑战。一方面，原始的Transformer自注意机制导致序列长度的时间和空间复杂度比较高。此外，相关研究通常使用多个Transformer层，导致难以承受的计算成本和内存占用，特别是当序列长度非常长时。为了解决这个问题，最近的许多研究集中在降低普通Transformer的计算复杂度上。另一方面，与cnn不同的是，transformer的性能严重依赖于数据集的规模。由于Transformer作为卷积运算不具有很强的感应偏置，因此在训练数据稀缺的情况下，Transformer的性能往往不如cnn。为了缓解这个问题，许多最先进的方法转向大规模数据集的预训练。然而，医学图像数据集普遍缺乏可用的训练样本，使得Transformer对医学图像的预训练变得不切实际。
  在本文中，作者利用3D CNN中的Transformer来有效地建模局部和全局特征来进行医学图像的体积分割。为了充分利用CNN和Transformer的优点，网络编码器首先利用3D CNN提取体积空间特征，同时对输入的3D图像进行下采样，有效地捕获局部3D上下文信息，生成紧凑的体积特征图。然后，每个卷被重塑为一个向量(即token)，并输入Transformer进行全局特征建模。最后，3D CNN解码器提取Transformer的特征嵌入和skip-connections的低层次特征表示，并进行递进向上采样来预测全分辨率分割图。作为一个混合的CNN-Transformer架构，作者的TransBTSV2可以在不进行任何前训练的情况下对医学图像进行精确分割。
  由于传统的基于Transformer的方法设计是沿着模型深度反复叠加Transformer层，因此Transformer的引入带来了相当大的计算开销。为了解决这个难题，受MobileNetV2中反向设计的启发，作者提出了一个新的见解，即追求更宽而不是更深的Transformer架构(即更喜欢更大的网络宽度而不是更深的网络深度)，以提高模型效率。在这种简单而有效的设计方案下，作者在整个流程中逐渐将模型宽度扩大两次，实现了反向瓶颈相似架构，与作者原来的TransBTS[1]相比，模型复杂度显著降低(参数减少了53.62%，FLOPs减少了27.63%)。由于只有一个Transformer层，作者的高效TransBTSV2实现了卓越的性能。
  通过对4个医学体积分割基准数据集进行综合实验，阐明将Transformer整合到3D CNN中，以释放两个网络的能力。值得注意的是，TransBTSV2是一个简单的、通用的3D网络(具有基本的3D卷积层和Transformer块)，没有任何复杂的结构和附加组件。TransBTSV2可以轻松地插入任何有效的技术，如多尺度特征融合，以提高性能。
  作者在本文中做了以下重要的扩展:
  1、作者提出了对Transformer模块进行重新设计的新见解，追求一种较浅但较宽的架构，而不是传统的较深但较窄的架构。该方法不仅大大降低了模型的复杂度，而且在医学图像分割中也取得了较好的性能。
  2、作者引入了一个**可变形的瓶颈模块(DBM)**来捕获更多的形状感知特征表示。利用所提出的DBM，作者的方法可以产生更精细的病灶区域分割结果。
  3、与MICCAI版本相比，作者在两个流行的医学图像数据集(lits 2017和KiTS 2019)上给出了更多的结果，以评估作者提出的TransBTSV2的泛化能力。在4个医疗体积分割基准数据集上的大量实验结果表明，TransBTSV2与之前的最先进的分割方法相比具有竞争力或更好的性能。

思路来源：
  ① U-Net和U-Net变体：
  由于上述医学图像分割的困难，对于先进的网络来说，多尺度特征融合和细粒度局部细节建模都是必不可少的。为了尽可能保留医学图像的细节信息，U-Net采用了带有跳跃连接的对称编解码结构，将下采样的特征映射逐渐恢复到原始大小，从而实现了医学图像的像素级密集预测。在接下来的研究中，基于U-Net的变体受到了广泛的关注，并被进一步应用到医学图像分割中。为了充分利用连续切片之间的局部信息，这也是体积分割的关键，Cicek等人通过实现3D操作，如3D convolution和3D max pooling，将U-Net从2D推广到3D，从而更好地进行医学体积分割。为了模拟远距离依赖关系，Ozan Oktay提出了一种新的医学成像注意门，它可以自动学会聚焦于不同形状和大小的目标结构。在ResNet残差结构的基础上，提出了残差UNets，以克服深层神经网络的训练困难。另外，还有研究，通过改进跳跃连接来丰富不同层次特征图的信息交互。

  ② Transformer-based Approaches：
   在大规模数据集上的预训练的帮助下，基于transformer的框架最近在各种视觉任务上达到了最先进的性能。Vision Transformer (ViT)将图像分割成固定大小的patch，并将这些patch之间的相关性建模为序列，得到了满意的结果。DeiT引入了一种知识蒸馏方法来训练ViT，以进一步提高其性能。DETR使用Transformer将对象检测视为集合预测任务，通过放弃生成锚点和非最大抑制的步骤，简化了对象检测过程。从序列到序列的角度重新思考语义分割，SETR利用Transformer作为编码器进行全局特征提取，并通过大规模的预训练获得卓越的性能。为了为分割任务构建高效的体系结构，SegFormer结合了一种新的分层结构的Transformer编码器和一个轻量级的多层感知器(MLP)解码器。此外，为了进一步实现高效的架构，减少冗余，Swin Transformer在ViT中去掉了全局的两两自注意机制，并在提出的局部窗口自注意机制和移位窗口方案的基础上，提出了一种新的Transformer主干，在各种基准上都达到了最先进的结果。
  为了明确地建模医学图像中体积的长距离依赖关系，近年来的研究主要集中在将Transformer引入医学图像分割任务中。TransUNet采用ViT作为编码器，并进行大规模的预训练，用于医学图像分割。TransFuse同时结合Transformer和CNN，以一种更浅显的方式有效地捕获全局依赖项和空间细节。此外，MedT提出了一种门控轴向注意模型，该模型在自注意模块中应用了额外的控制机制来扩展现有的体系结构。MISSFormer用增强Transformer块重新设计了前馈网络，以增强长期依赖关系并补充本地上下文。同时，提出了一种增强型Transformer上下文桥，对分级Transformer编码器生成的多尺度特征的长期依赖关系和局部上下文进行建模。TransAttUnet在Transformer中集成了一种新的自意识注意模块，并提出了全局空间注意模块，以有效地学习编码器特征之间的非局部交互。与之前提出的这些2D方法不同，作者的重点是研究如何有效地将Transformer整合到流行的3D cnn中，以释放这两个网络的潜力。作者的TransBTSV2不仅通过复杂的设计超越了SOTA的结果，而且是具有卓越性能的高效模型。

模型方法:

  图2说明提出的用于自动医学图像分割的TransBTSV2。使用修改过的3D CNN编码器来捕获局部信息，并利用Transformer编码器从全局视图建模远程依赖关系。将上采样层和卷积层叠加，逐步产生高分辨率分割结果。在整个架构中，作者将模型宽度扩展了两倍(如图中用红色突出显示的)。TransBTSV2具有类似于Transformer架构的倒置瓶颈(如图中黄色阴影部分所示)和拟议的DBM，因此TransBTSV2是一个具有卓越性能和低模型复杂性的高效架构。

可变形的瓶颈模块：
  由于医学图像中病变区域形状不规则且种类繁多，精确识别病变区域具有很大的挑战性。为了在医学图像分割中获得更好的性能，网络需要捕捉细粒度的、形状感知的局部细节。由于CNN基本模块的固定几何结构，CNN固有地局限于模型的不规则形状变形。为了解决这一问题，设计了变形瓶颈模块，进一步从不规则形状病变区域捕获形状感知特征。由于编码器的特征映射对几何信息更加敏感，对目标区域的识别至关重要，因此提出的DBMs被插入到每个跳过连接中。每个DBM由两个1×1×1卷积、一个3×3×3可变形卷积和传统残差连接组成。为了最大限度地减少提出的DBM带来的计算开销，作者部署了两个1 × 1 × 1卷积(即图3所示的Reduction和Restoration layer)来降低和恢复信道维数，同时利用3D可变形卷积来捕获形状感知的局部信息。数据库管理系统在每个跳跃连接可以有效地从编码器特性中学习体积空间偏移量，并适应各种分割目标的转换。因此，从DBMs输出的特征映射提供了更多形状感知的局部细节，作者提出的TransBTSV2可以生成更精细的损伤区域分割结果。

实验与结果：
  数据与评价指标：
  在四个基准数据集上评估了作者提出的方法用于医学体积分割。
  1) BraTS2019和BraTS2020:实验中使用的第一个3D MRI数据集由脑瘤分割(BraTS) 2019挑战提供。它包含了335例用于培训的患者和125例用于验证的患者。每个样本都由四种形式的大脑MRI扫描组成。每个形态都有一个240 × 240 × 155的体积，已经对齐到相同的空间。标签包括4类:背景(标签0)、坏死和非增强肿瘤(标签1)、瘤周水肿(标签2)和gd增强肿瘤(标签4)。在线评估服务器评估的验证集性能用于验证所提方法的有效性。骰子的分割精度测量分数和豪斯多夫距离(95%)指标提高肿瘤区域(ET、标签1),肿瘤区域的核心(TC、标签1和4),和整个肿瘤区域(WT、标签1、2和4)。第二个3 d核磁共振数据集是由大脑肿瘤分割2020年挑战[39]-[41]。培训369例，验证125例，测试166例。除了数据集中的样本数量外，关于这两个MRI数据集的其他信息是相同的。
  2) LiTS2017:第三个医学图像数据集由肝脏肿瘤分割(LiTS) 2017 challenge[42]提供。LiTS数据集包含131个和70个对比增强3D腹部计算机断层扫描样本，分别包含用于训练和测试的肝脏和肝脏肿瘤的标注。切片数为75 - 987，每个CT图像和标签的空间大小为512×512。标签包含3类:背景(标签0)，肝脏(标签1)，肝脏肿瘤(标签2)。作者采用每例Dice评分和Dice整体评分作为评价指标。Dice per case score指的是每个卷的平均Dice得分，Dice global score是将整个数据集合并到一个数据集中评估的Dice得分。
3) KiTS2019:第四医学图像数据集由肾肿瘤分割(KiTS) 2019 Challenge[43]提供。试剂盒2019数据集为300名患者提供具有高质量标注体素标签的多期3D ct数据。随机选取210例患者作为训练集，其余90例为算法评价测试集。ground truth包含3个类别:背景(标签0)、肾脏(标签1)、肾脏肿瘤(标签2)。对于作者的方法的评估，使用与KiTS 2019挑战相同的3个评估指标。Kidney dice表示同时考虑肾脏和肿瘤作为前景时的分割性能，而tumor dice则考虑除肿瘤以外的一切作为背景。复合dice是肾脏dice和肿瘤dice的平均值。

评价指标:
  使用Dice相似系数(Dice)和Hausdorff距离(HD)作为本研究的性能指标。Dice测量分割预测和地面真实之间的重叠，HD通过计算预测边界和地面真实之间的最大距离来评估分割边界的质量。

实现细节:
  TransBTSV2基于Pytorch实现，并使用NVIDIA Titan RTX图形处理器进行训练。在四个基准数据集上，TransBTSV2从零开始训练，批处理大小为16。在BraTS 2019和BraTS 2020数据集上，扩展比E设置为1.5。采用Adam优化器对6000个时代的模型进行训练。在训练过程中采用60个epoch的热身策略，初始学习速率设置为0.00008，并采用余弦学习速率衰减计划。在训练阶段采用随机裁剪、随机镜像翻转和随机强度移动作为数据增强。采用softmax Dice loss训练网络，采用L2 Norm进行模型正则化，权值衰减率为10−5。在LiTS2017和KiTS2019数据集上，扩展比E设置为2。用Adam优化器分别对6000个epoch和12000个epoch进行训练。在训练过程中分别采用80个epoch和160个epoch的热身策略，初始学习速率设置为0.002，采用余弦学习速率衰减计划。

  为了公平地评估作者提出的TransBTSV2的性能，作者首先在BraTS 2019验证集上进行了实验，并将TransBTSV2与之前的最新方法进行了比较。定量结果如表一所示，TransBTSV2在ET、WT、TC上的Dice得分分别为80.24%、90.42%、84.87%，Hausdorff距离分别为3.696mm、5.432mm、5.473mm，与表一所示的以往SOTA方法的结果相当或更高。与以前的SOTA方法相比，在分割方面已经取得了相当大的改进。此外，可以看到，TransBTSV2通过作者重新设计的Transformer和DBM，获得了ET和TC的Hausdorff距离的绝对优势。与没有Transformer作为编码器一部分的3D U-Net[26]相比，TransBTSV2在性能和效率上都有很大的优势，有了显著的改进。这清楚地揭示了利用Transformer建模全局关系的好处。

          2019年验证集的表现比较
定性分析：

  图4所示。MRI脑肿瘤分割结果的视觉比较。蓝色区域表示增强肿瘤，红色区域表示非增强肿瘤，绿色区域表示瘤周水肿。
  还对BraTS 2020验证集上的TransBTSV2进行了评估，结果见表二。TransBTSV2在ET、WT、TC上的Dice得分分别为79.63%、90.56%、84.50%，Hausdorff距离分别为12.522mm、4.272mm、5.560mm。与3D U-Net[26]、V -Net[45]和Residual U-Net[9]相比，TransBTSV2在这两个指标上都表现出了很大的优势，并有明显的改进。此外，在BraTS 2020验证集上，TransBTSV2实现了比以前的SOTA方法更具竞争力或更好的性能，证明了作者精心设计的混合架构的潜力。值得一提的是，与TransBTS相比，TransBTSV2在BraTS 2019和BraTS 2020两项指标上都取得了巨大的进步。

  为了评价TransBTSV2的性能和泛化能力，作者还在litit 2017数据集上进行了实验。各评价指标的绩效比较见表三。可以发现，与SOTA方法相比，TransBTSV2在广泛的度量标准中再次获得了具有竞争力或更好的性能。特别是对于病变区域的Dice指标，作者的TransBTSV2取得了相当大的改善，得益于DBM和改进的Transformer的开发。为了更直观的比较，不同方法的分割结果如图5所示。很明显，与U-Net[26]、V -Net[45]和Attention U-Net[12]相比，作者的分割结果更接近地面真相。TransBTSV2结合了transformer和convolutions的优点，可以更精确的分割肝脏肿瘤，产生更精细的分割结果。

  为了进一步验证作者的方法的有效性和鲁棒性，作者在KiTS 2019数据集上进行实验，该数据集提供的肾脏和病变的图像和标签种类多、复杂性高。表四为作者的方法与以往SOTA方法在肿瘤和肾脏分割方面的性能比较。与triplstage 3D U-Net[63]相比，一个比作者的TransBTSV2模型复杂得多的三级体系结构，可以看到在复合骰子度量方面取得了显著的改进(0.43%)。从表四可以看出，TransBTSV2进一步提高了对肾脏和肿瘤的分割精度，证明了精心设计的混合架构的有效性。为了分析分割结果的质量，肾脏和肿瘤的三维显示如图6所示。可以看出，利用Transformer和提出的DBM模型可以模拟长时间的依赖关系和不规则形状的变形，作者的方法可以很好地分割小尺度肿瘤和相对较大的肾脏。

  TransBTS是一个中等大小的模型，具有32.99M参数和333G FLOPs。通过本文提出的改进架构设计，TransBTSV2只有15.30M参数和241G FLOPs，而在BraTS2019验证集上ET、WT、TC的Dice得分分别为80.24%、90.42%、84.87%，Hausdorff距离分别为3.696mm、5.432mm、5.473mm。值得注意的是，与作者原来的TransBTS[1]相比，追求更宽而不是更深的Transformer可以大大降低复杂度(参数减少53.62%，FLOPs减少27.63%)，但同时显著提高了模型性能。与具有16.21M参数和1670G FLOPs的3D U-Net[26]相比，TransBTSV2在计算复杂度和模型性能方面都有很大的优势。

消融实验：
  作者进行了大量的消融实验来验证TransBTSV2的有效性，并基于BraTS 2019验证集的评估来证明其设计的合理性。因为挑战没有提供测试集，消融实验是在官方验证集上进行的，以进行公平的比较。如表V所示，与作者的基线(修改后的3D U-Net)相比，混合CNN-Transformer架构在Dice分数和Hausdorff距离两个指标上都显示出极大的优越性。这清楚地揭示了利用Transformer建模全局关系的好处。此外，利用所提出的DBM使得TransBTSV2能够更有效地模拟病变区域的不规则形状变形。

  作者还对Transformer中空间片自注意的不同设计进行了深入的研究。联合版本同时对空间-片自注意进行集体计算，分割版本将联合注意分别解耦为空间注意和片注意(即spatial Transformer和slice Transformer)，仅计算空间维度上的自注意。级联和并联表示上述两种解耦Transformer(即空间Transformer和切片Transformer)的相对位置。如表六所示，联合版本在空间维度和切片维度同时计算自注意，其性能优于其他设计(拆分版本和仅空间版本)。联合版在ET和TC的Dice分数上有了很大的改进。虽然分割版本和仅空间版本可以降低大量的内存成本，但由于在单个维度(即空间或片)中标记的数量(即序列长度)有限，它们不能有效地建模体积空间特征之间的显式长程依赖关系，导致模型性能的急剧下降。

  在本文中，作者提出了一种新的分割框架，将Transformer集成到3D CNN中，用于医学图像的体积分割。最终的架构TransBTSV2不仅继承了3D CNN在建模本地上下文信息方面的优势，而且还利用Transformer学习全局语义相关性。TransBTSV2作为一种混合的CNNTransformer架构，可以在不进行任何前训练的情况下对医学图像进行精确的分割，具有与cnn一样的强烈的诱导偏差。在4个医学图像数据集上的实验结果表明，TransBTSV2不仅高效，而且在性能上优于或可与之前的先进方法相媲美。