[深度学习论文笔记]UCTransNet:从transformer的通道角度重新思考U-Net中的跳跃连接

UCTransNet: Rethinking the Skip Connections in U-Net from a Channel-wise Perspective with Transformer
UCTransNet:从transformer的通道角度重新思考U-Net中的跳跃连接

Published : AAAI 2022
论文：https://arxiv.org/abs/2109.04335
代码：https://github.com/mcgregorwwww/uctransnet

摘要：
  大多数最新的语义分割方法采用了一个带有编码器-解码器结构的U-Net框架。对于具有简单跳过连接方案的U-Net来说，对全局多尺度上下文建模仍然具有挑战性:
  1)由于编码器和解码器级的不兼容特征集的问题，并非每个跳过连接设置都是有效的，甚至一些跳过连接也会对分割产生负面影响性能；
  2)在某些数据集上，原U-Net比没有任何跳转连接的U-Net更差。基于作者的发现，作者提出了一个新的细分框架，叫做UC- TransNet(在U-Net中有一个建议的CTrans模块)，从通道的角度来看，有注意机制。具体来说，CTrans(通道转换器)模块是U-Net跳跃连接的一种替代，它由一个进行多尺度通道交叉融合的子模块(命名为CCT)和一个引导融合的多尺度通道交叉信息的子模块(命名为CCA)组成，以有效地连接到解码器特征，从而消除歧义。因此，所提出的由CCT和CCA组成的连接能够代替原始的跳过连接来解决语义鸿沟，从而实现精确的医学图像自动分割。实验结果表明，对于跨不同数据集和涉及transformer或U形框架的传统体系结构的语义分割，作者的UCTransNet产生了更精确的分割性能，并实现了对现有技术的一致改进。

  图1:本文提出的UCTransNet (d)跳过连接方案与其他模型的比较。虚线表示跳过连接。

问题动机：
  医学成像被认为是帮助医生评估疾病和优化预防和控制措施的重要技术。医学图像中目标对象的分割和随后的定量评估为病理分析提供了有价值的信息，并且对治疗策略的规划、疾病进展的监测和患者结局的预测很重要。最近的方法语义分割通常依赖于卷积编码器-解码器架构，其中编码器生成低分辨率图像特征和解码器将特征上采样到具有每像素类分数的分割图中。U-Net是用于医学图像分割的最广泛使用的编码器-解码器网络结构，因为编码器捕获低级和高级特征，而解码器组合语义特征来构造最终结果。跳过连接有助于纠正在池操作期间丢失的空间信息，从而通过编码-解码过程恢复全部空间分辨率。为了研究它，作者对U-Net进行了深入研究，并根据作者对多个数据集的分析观察到了几个主要的限制。作者发现，对于具有简单跳跃连接方案的U-Net来说，建模全局多尺度上下文以辅助解码过程而不考虑语义间隙仍然是一个挑战。为了实现精确的医学图像分割，有必要寻找一种有效的融合特征的方法。对于U-Net的扩展，本质上有两个关键问题:编码器中的哪些特征层连接到解码器，用于通过对多尺度特征进行聚集来对全局上下文进行建模，以及如何有效地将这些特征与可能的语义间隙相融合，而不是简单地进行连接？存在两个语义鸿沟:多尺度编码器特征之间的语义鸿沟以及编码器和解码器级之间的语义鸿沟，限制了分割性能。为了克服上述限制，最近引入了许多方法来缓解融合这两组不兼容特征时的差异。一种方法是用嵌套密集跳过路径直接替换普通跳过连接，用于医学图像分割。最有代表性的方法是UNet++ 其通过引入具有一系列卷积的密集连通性来缩小编码器和解码器子网络之间的语义差距，并实现更好的分段性能。这是对U-Net中只需要融合相同比例的特征图的限制性跳跃连接的改进。另一种方法侧重于通过在从编码器级传播的特征上引入额外的非线性变换来加强跳跃连接，这应当解释或在某种程度上平衡可能的语义间隙。
  尽管取得了良好的成绩，上述两个网络仍然无法从整体上有效地探索足够的信息。在医学图像分割中，捕捉多尺度特征对于解决复杂的尺度变化至关重要。在这些重要问题的驱动下，出现了一个问题:如何通过有效捕获非局部语义依赖，通过多尺度信道信息融合来充分弥合编码器和解码器之间的语义鸿沟。在本文中，作者重新思考了跳跃连接设计，并提出了一种替代方法来更好地连接编码器和解码器级之间的特性。不同的通道通常关注不同的语义模式，自适应地融合足够的通道特征有利于复杂医学图像的分割。为此，作者提出了一种以U-Net为网络主体结构的端到端深度学习网络UCTransNet。更具体地说，作者首先提出了通道式交叉融合转换器(CCT)，从通道式的角度融合多尺度上下文和交叉关注。该方法的目标是捕获局部跨通道交互，通过协作学习而不是独立连接来有效融合可能存在尺度语义差距的多尺度通道特征。另一方面，为了解决语义层次不一致的问题，作者提出了另一种融合多尺度特征和解码阶段特征的通道交叉注意模型。这两个交叉注意模块都被称为CTrans (Channel Transformer)，它可以通过探索多尺度的全局上下文来建立编码器和解码器之间的关联，并替换原有的跳过连接来解决语义鸿沟，从而提高分割性能。这两个模块都可以很容易地嵌入并应用于医学图像分割任务中的U形网络。大量实验表明，与U-Net相比，UCTransNet在GlaS、MoNuSeg和Synapse数据集上分别获得了4.05% Dice、7.98% Dice和9.00% Dice的绝对增益，从而极大地改善了传统的分割流水线。此外，作者做了一个彻底的分析来研究特征交互是如何工作的。此外，之前的工作已经结合了transformer和U-Net来显式地模拟远程空间依赖性. 结果表明，通道式融合transformer方案通常比用变压器代替卷积运算的方法性能更好。作者认为UCTransNet可以作为医学图像分割的强跳跃连接方案。
贡献有三点：
  1)作者的研究是第一个在多个数据集上充分探索U-Net中跳过连接的潜在弱点并发现独立简单复制是不合适的工作。
  2)作者提出了一个新的视角来促进语义分割性能，即通过更有效的特征融合与多尺度通道式交叉关注来捕捉更复杂的通道式相关性，从而缩小低级和高级特征之间的语义和分辨率差距。
  3)作者的方法是一种更合适的U-Net和Transformer的组合，具有更少的计算成本和更高的性能。与其他最先进的分割方法相比，实验结果在所有三个公共数据集上都表现出更好的性能。

图2:拟议的UCTransNet的示意图。作者将原来的跳过连接替换为由两个组件组成的CTrans:通道交叉融合变压器(CCT)和通道交叉注意(CCA)。

思路来源：
  ① 用于医学图像分割的transformer：
  最近，Vision Transformer (ViT)通过直接将具有全球自我关注功能的Transformer应用于全尺寸图像，实现了ImageNet分类的最新技术。由于《transformer》在许多计算机视觉领域的成功，一种新的医学图像分割范式最近得到了发展。TransUNet (Chen et al. 2021)是第一个基于Transformer的医学图像分割框架。Valanarasu等人提出了一种门控轴向注意模型- medt (Valanarasu等人，2021)，以克服医学成像中数据样本数量少的问题。在实现最先进性能的Swin Transformer的激励下，Swin- unet提出了第一个纯基于Transformer的U型架构，该架构引入了Swin Transformer来取代U-Net中的卷积块。但上述方法主要针对的是卷积运算的缺陷，而不是U-Net本身，可能会造成结构冗余，计算成本过高。

  ② Unet中的跳跃连接：
  跳过连接机制首先在UNet中提出，其设计目的是弥合编码器和解码器之间的语义鸿沟，并已被证明在恢复目标对象的细粒度细节方面是有效的。U-Net的普及后,提出了许多新颖的模型如UNet + +,attention U-Net, DenseUNet, R2U-Net,和UNet 3 +,是专门设计用于医学图像分割,实现表达性能。Zhou等人认为，来自编码器和解码器网络的同尺度特征映射在语义上是不同的，因此设计了一种名为UNet++的嵌套结构，该结构捕获多尺度特征，进一步弥补了这一差距。Attention UNet提出了交叉注意模块，该模块利用粗尺度特征作为门控信号，消除跳跃连接中不相关和噪声响应的歧义。MultiResUNet观察到在同一层次上跳过的编码器特征和解码器特征之间可能存在语义差距，因此他们引入了带有残差结构的Res Path来改善跳过连接(见图1)。
这些方法假设每个跳过连接都有相等的贡献，但是在下一节中，我们将展示所有跳过连接的贡献是不同的，有些甚至可能会损害最终的性能。

  ③ 跳跃连接分析：
  发现1:没有任何跳过连接的U-net甚至比原来的U-net更好。比较图3的结果，作者发现，在MoNuSeg数据集上，对于几乎所有的度量，“U-Net-none”算法的性能最差。然而，“U-Net-none”，尽管没有任何约束，仍然在GlaS数据集上与“U-Net-all”相比取得了非常有竞争力的性能。它证明了跳过连接并不总是有利于分割。
  发现2:虽然UNet-all比UNetnone执行得更好，但并不是所有带有简单复制的跳过连接都对分段有用。每个跳过连接的贡献是不同的。作者发现，在MoNuSeg数据集上，对于Dice和IOU，每个跳过连接的性能范围分别为[67.5%，76.44%]和[52.2%，62.73%]。对于不同的单跳线连接，冲击变化较大。此外，由于编码器和解码器阶段的特征集不兼容的问题，一些跳过连接对分割性能有负面影响。例如，在GlaS数据集上的Dice和IOU方面，l1的表现比UNet-none差。结果并没有表明，编码器阶段的许多特性是没有信息的。其原因可能是简单复制不适用于特征融合。
  发现3:对于不同的数据集，skip贡献的最佳组合是不同的，这取决于目标病灶的规模和外观。作者进行了几个消融实验来探索最好的侧输出设置。由于空间有限，作者忽略了两个跳过连接的组合。可以看到，跳过连接并没有获得更好的性能。带l4的模型在MoNuSeg数据集上表现最好，而让我们惊讶的是，只有一个跳过连接的l3在GlaS数据集上表现最好。这些观察结果表明，对于不同的数据集，最佳组合是不同的，这进一步证实了引入更合适的动作过程来进行特征融合的必要性，而不是简单的连接。

模型方法:
  用于医学图像分割的UCTransNet

   图2展示了UCTransNet框架的概述。据作者所知，目前基于transformer的分割方法主要集中在对U-Net编码器的改进上，因为它具有捕获远程信息的优势。这些方法，如TransUNet 或TransFuse，以一种简单的方式将Transformer与U-Net混合，即将Transformer模块插入编码器或融合两个独立的分支。然而，作者认为当前的U-Net模型的潜在限制是跳跃连接的问题，而不是原始的U-Net的编码器问题，这足以满足大多数任务。正如在跳过连接分析一节中提到的，作者观察到编码器的特征与解码器的特征不一致，即在某些情况下，由于浅层编解码器与浅层编解码器之间存在语义差异，语义信息较少的浅层特征通过简单跳转连接可能会影响最终性能。受此启发，作者构建了UCTransNet框架，在普通的U-Net编码器和解码器之间设计了一个通道式Transformer模块，以更好地融合编码器特性，减少语义差距。具体而言，作者提出了一种通道Transformer (CTrans)来替代U-Net中的跳过连接，它由两个模块组成:CCT (Channel-wise Cross Fusion Transformer)用于多尺度编码器特征融合，CCA (Channel-wise Cross - Attention)用于译码器特征和增强的CCT特征融合。

  图4:原始的self-attention(a)和作者建议的通道交叉注意(b)之间的比较。

  图5:多头交叉注意

  CCT:用于编码器特征转换的通道式交叉融合Transformer：
  为了解决之前提到的跳过连接问题，作者提出了一种新的基于通道的交叉融合Transformer (CCT)来融合编码器的多尺度特性和Transformer长依赖建模的优势。CCT模块包括三个步骤:多尺度特征嵌入、多头通道交叉注意和多层感知器(MLP)。

实验与结果：

  表1:GlaS和MoNuSeg数据集的3倍5倍交叉验证结果。dice和借据采用“平均±标准”格式。符号∗表示我们的方法在该分数上明显优于其他方法(使用学生t检验，水平为0.05)。

  表2:与Synapse数据集上最先进的分割方法的比较。为简单起见，’ R50-U-Net ‘和’ R50-AttUNet '分别表示以ResNet-50为骨干的U-Net和Attention U-Net。

数据集：
  作者使用Gland segmentation， MoNuSeg和Synapse多器官分割数据集来评估我们的方法。Gland segmentation dataset (GlaS)有85张用于训练的图像和80张用于测试的图像。MoNuSeg数据集有30张图片用于训练，14张用于测试。Synapse腹部CT共扫描30次，共扫描8个腹部器官(主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、脾、胃)，轴向CT图像共3779张。接下来，作者使用18个训练案例(2212个轴向切片)和12个案例进行验证。

  图6:GlaS和MoNuSeg数据集的定性比较

  图7:Synapse数据集上的定性比较

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

实现细节:
  作者使用PyTorch在单个NVIDIA A40 GPU卡上实现了我们的模型，内存为48 GB。为了避免过拟合，作者还进行了水平翻转、垂直翻转和随机旋转两种在线数据增广。没有使用任何预先训练的权值来训练提议的UCTransNet。对于GlaS和MoNuSeg，将批量大小设置为4，而对于Synapse，将其设置为24。三个数据集的输入分辨率和patch size P均设为224 × 224和16。使用Adam优化器来训练模型，初始学习率设置为0.001。还利用交叉熵损失和dice损失作为损失函数来训练我们的网络。为了使小数据集的结果更具有说服力，进行了3次5倍交叉验证(共15个CV)，并获得了平均结果和std。统计检验表明，作者的方法明显优于可比方法。对于GlaS和MoNuSeg数据集，作者使用dice系数(dice)和联合交叉(IoU)作为评估指标，而对于Synapse数据集，作者报告骰子和HausdorffDistance(HD)。注意，使用相同的设置和loss函数训练所有的基线。
  为了演示所提出的UCTransNet的整体分割性能，作者将其与其他最先进的方法进行比较。将UCTransNet与两种方法进行综合评价的比较，包括三种基于UNet的方法:UNet++、Attention U-Net、MultiResUNet和三种最先进的基于变压器的分割方法，包括TransUNet、MedT和SwinUnet。为了进行公平的比较，实验中使用了它们最初发布的代码和发布的设置。本文还介绍了两种优化UCTransNet模型的策略。1)联合训练:优化U-Net中的卷积和CTrans参数，优化两个通道方向的交叉注意参数，同时优化一个损耗;2)训练。首先训练一个U-Net，然后用同样的数据进一步训练UCTransNet中的参数。
  实验结果报告在表1中，最好的结果被加粗显示。表1显示作者的方法比现有技术有持续的改进。在表2中，可以得出类似的观察和结论，这再次验证了UCTransNet优于其他所有公司。此外，该预训练方案不仅具有更快的收敛速度，而且在MoNuSeg数据集上的性能优于竞争方法，甚至优于联合学习方案。这些观察结果表明，这两个拟议的模块可以纳入预先训练的U-Net模型，以改善分割性能。提供了参数数字和gflops，这表明我们的模型在有效性和效率之间取得了很好的平衡。

  表3:GlaS和MoNuSeg数据集上的消融实验。“CCT”表示建议的信道变压器，“CCA”表示基于信道的交叉注意。最好的结果用黑体字表示。

  图8:对GlaS数据集和MoNuSeg数据集上的查询和键数的消融研究

消融实验：
  如表3所示，‘Base+CCT+CCA’在所有数据集上普遍优于其他基线，这说明了两个模块结合的有效性。我们的研究结果揭示了多尺度多通道特征融合在编码器-解码器框架中提高分割性能的重要性。
  前面的实验表明，我们的模型中的CCT模块能够有效地增强跳跃连接。在之前的实验中，所有编码器级别的多尺度特征都参与到CCT模块中，因此查询次数为4，关键是由4个尺度特征组成的串接表示。
  对编码器和解码器之间的跳过连接数量进行了一系列的实验，如图8所示。注意，关键向量是固定的，它仍然由四个尺度特征组成。随着跳过连接数量的增加，我们观察到一致的改进。这一观察结果表明，不同编码器级别学习的多尺度特征是有用的，这验证了作者的动机。
  从图8中可以看出，性能随着特征尺度的增加而提高，达到4个尺度，说明通道越多越能准确捕捉节点特征，说明将特征尺度越大转化为查询越好

  图9:GlaS数据集(a)和MoNuSeg数据集(b)的相似度矩阵。’ K1 ‘表示与key中连接的’ Q1 '相同的特征。

  为了对UCTransNet进行全面评估，作者在图9中可视化了CCT模块中的交叉注意分布。研究哪一级别的编码器具有更强的自信相关性，对分割更重要，也是一个有趣的问题。可以看出，“K2”和“K3”分别与GlaS和MoNuSeg数据集上的其他编码器级别有更自信的相关性。研究结果与图3中U-Net中箕斗连接分析结果一致。它解释了为什么’ L3 ‘和’ L2 '在GlaS和MoNuSeg数据集上分别获得更好的性能。这意味着有必要开发一种多尺度特征融合来解决语义差距问题，这也验证了作者构建全局多尺度通道特征融合模型来有效捕获非局部语义依赖关系的动机。

总结：
  医学图像的准确、自动分割是临床诊断和分析的关键步骤。在这项工作中，作者从通道的角度介绍了一个通道transformer分割网络(UCTransNet)，以提供精确和可靠的医学图像自动分割。该方法通过端到端结合多尺度通道交叉融合transformer(CCT)、循环神经网络和通道交叉注意(CCA)的优势，显著改善了当前医学图像在多基准数据集上的分割效果。通过深入的分析和实证，transformer展示了UCTransNet模型的优势。它成功地缩小了语义鸿沟，充分利用了编码阶段的多尺度特征。