通过特征蒸馏与迭代学习改进 UNet 的图像分割能力！

未来先知

发布于 2024-08-08 11:17:11

1980

发布于 2024-08-08 11:17:11

文章被收录于专栏：未来先知

自从UNet被引入以来，它一直在引领各种医学图像分割任务。尽管许多后续研究也致力于提高标准UNet的性能，但很少有研究深入分析UNet在医学图像分割中潜在的感兴趣模式。在本文中，作者探讨了在UNet中学习到的模式，并观察到了可能影响其性能的两个重要因素： (i)由不对称监督引起的迭代特征学习； (ii)特征图中的特征冗余。为此，作者建议平衡编码器和解码器之间的监督，并减少UNet中的冗余信息。具体来说，作者使用包含最多语义信息（即解码器的最后一层）的特征图来为其他块提供额外的监督，并通过利用特征蒸馏来减少特征冗余。所提出的方法可以轻松地以即插即用的方式集成到现有的UNet架构中，计算成本几乎可以忽略不计。实验结果表明，所提出的方法一致地提高了标准UNet在四个医学图像分割数据集上的性能。代码可在https://github.com/ChongQingNoSubway/SelfReg-UNet获取。

1 Introduction

医学图像分割是计算机辅助诊断和图像引导系统中的关键应用。近年来，深度学习已成为这一领域的主要方法，这主要归功于UNet[18]的里程碑式贡献。UNet通过利用编码器将语义信息投射到低级特征，以及解码器逐步将语义特征上采样到分割 Mask ，定义了一个通用的分割网络架构。其许多后续工作在卷积神经网络（CNN）的背景下扩展了这一想法。最近，许多研究引入了视觉 Transformer ，通过使用自注意力机制来解决CNN的限制。尽管ViT具有较大的感受野，并能捕获不同图像块之间的长距离依赖，但由于缺乏局部性，它难以保留细粒度的局部上下文。为了缓解这个问题，一些方法[8, 17]在UNet设计中引入了CNN和ViT之间的桥梁（即混合模型）。请注意，这些方法也带来了更多的计算复杂性和参数数量。过参数化是深度学习中的常见问题，通常导致特征冗余和特征表示不佳[6, 12, 13]。然而，这个问题在当前的医学分割模型中尚未正式研究或考虑。

除了上述方法，一些工作集中在优化UNet的结构。沿着这一思路，Att-Unet[15]提出了一种基于注意力的跳跃连接来抑制不相关的特征。Unet++[30]用嵌套密集跳跃路径替换了标准的跳跃连接（即拼接/加法）。UCTransNet[23]彻底分析了不同跳跃连接的影响，并提出了通道 Transformer 以替换传统的跳跃连接。这些方法只通过操作跳跃连接来探究编码器到解码器之间的信息流。然而，它们都没有探索如何有效地通过解码器中学到的特征来通知编码器，而作者的研究揭示了这个信息流在UNet中值得更多关注。这主要是因为解码器比编码器接收到更多的监督，这提供了一种自然的方式来过滤掉不相关的信息。

在本文中，作者对两种代表性的UNet（即标准的Unet[18]和SwinUnet[1]）进行了实证研究。

作者的分析揭示了两个关键发现：

（i）特征通道中存在冗余特征，特征图中的浅通道比深通道表现出更多的多样性；

（ii）UNet中编码器和解码器之间的不对称监督导致语义损失。这种现象与其他计算机视觉任务[28, 25]中观察到的趋势不同，在这些任务中，深特征表现出更具辨别性，并能更好地定位感兴趣的目标。为了减轻这些问题，作者引入了语义一致性正则化和内部特征蒸馏，分别解决不对称监督和特征冗余导致的语义损失。这涉及到使用更准确的语义来监督其他块，并在特征图的浅通道到深通道之间进行信息蒸馏。

贡献：

作者的探索揭示了UNet中的不对称监督和特征冗余，为未来模型设计指出了一个新颖的方向。
作者提出了一种正交方法，通过提出对称监督正则化机制和利用特征蒸馏，帮助UNet丢弃不相关信息，更好地保留语义。
所提出的方法可以无缝集成到现有的UNet框架中（例如，CNN-UNet和ViT-UNet），以最小的额外成本提供性能提升。

2 Method

Preliminary

在本文中，作者以文献[18]中定义的标准Unet为例，其深度为5。为了普遍性，作者在研究中考虑了基于CNN的Unet [18]和基于ViT的Unet [1]的统一Unet结构，该结构包括输入和输出投影块、4个编码器和解码器块以及一个瓶颈（见图1(a)）。为此，作者首先将ViT-Unet中的块嵌入块定义为输入投影块与第一个编码器块的组合。然后，作者将ViT-UNet中的最后一个上采样块（即SwinUnet [1]中的最后一个块扩展）定义为最后一个解码器块。所定义的UNet的每个编码器/解码器块包含个连续的卷积/ Transformer 块。输出投影块是一个卷积层，它将最后一个特征图映射到一个分割 Mask 。为了表示方便，作者使用和来表示 -th 编码器/解码器块的 -th 层和瓶颈，分别。相应地，按照从编码器到解码器的顺序，特征图表示为（，,...,）。

Analysis on features learned in a Unet

图1：（a）Unet结构。（b）基于ViT/CNN的UNet中，每个编码器和解码器对应的注意力图。（更多示例参见补充附录A）(c) 基于ViT/CNN的Unet在浅层（左）和深层通道（右）之间的特征相似性矩阵。

作者通过采用两种常用的技术来分析在UNets中学习到的特征：

（i）梯度加权的类激活映射（Grad-CAM）[19]；

（ii）特征图中的相似性分析。UNet中的不对称监督。

作者从图1(b)中观察到了两个有趣的现象：

（i）学习模式在编码器和解码器之间存在不对称性。解码器可以近似定位一些 GT 分割区域，而编码器倾向于捕捉不相关信息（，），将兴趣模式分散到边界。

（ii）在解码器中，位于更末端位置的块（）对 GT 分割有了准确的理解。同时，块（,）学习了不相关信息。这主要源于每个块接收到的监督信号强度不同。当从直接由 GT 值监督的输出投影回溯时，监督信号逐渐减弱。它导致了语义损失，使得编码器中的一些块（例如,,,B）甚至激活了与分割无关的区域。

UNet中的冗余特征。实证研究表明，过参数化的CNN/ViT模型倾向于学习冗余特征，导致视觉概念不佳。以的输出为例，作者在通道维度上计算了浅层和深层 Level 的特征相似性矩阵。

如图1(c)所示，作者在基于ViT/CNN的UNets中观察到了两个现象：

（i）深层中普遍存在特征冗余，高相似性矩阵表明跨通道学习到相似特征。

（ii）浅层表现出显著的多样性，低相似性矩阵证明了这一点。

Unet中存在的过参数化，与其他网络类似，是这些现象的根本原因。这些冗余特征通常伴随着与任务无关的视觉特征，导致性能下降和不必要的计算开销。

Solutions

语义一致性正则化。 已有众多研究通过采用诸如知识蒸馏和特征对齐等策略来处理自然图像中语义信息的丢失，目的是利用准确的特征来指导那些信息量较少的特征[2, 13, 9, 5, 27]。受到这些在自然图像方面工作的启发，作者提出使用包含最多语义信息（即观察到的）的特征图来为UNet中的其余块提供额外的监督。因此，作者提出了一种通用的范式，称为语义一致性正则化（SCR），以平衡编码器和解码器之间的监督。为了简单起见，作者使用了在[2, 13]中提出的特征蒸馏机制来阐述作者的想法，因其普及性和简便性（以均方损失的形式给出）；而其他知识蒸馏方法可以作为替代（以KL散度形式给出）。为此，作者定义了所提出的SCR为

其中是位于最后一个解码器块（）的特征图，而是除了之外，位于第m个块的 ith 层的所有特征图（，,...，,,）。为了在通道和空间维度上对齐特征，作者采用了如图2（a）所示的均值池化和随机通道选择操作（RSC）。值得注意的是，通道选择并没有引入额外的模块[9, 5, 27]，从而减少了计算和语义冲突。范数被用作距离度量。

内部特征蒸馏。 为了解决特征冗余问题，一些与模型滤波器剪枝领域相关的通道缩减方法已经被提出，它们利用范数惩罚来对通道显著性引入稀疏先验[29, 24, 14]。受到这一启发，作者使用范数从浅层（上半部通道特征）到深层（下半部通道特征）进行信息蒸馏，这指导了深层特征学习有用的上下文信息。它可以表述为：

其中表示位于第m个块的第i层的所有特征图（，,..., ,,），是深层通道特征，是浅层通道特征。如图2(b)所示，作者将通道划分为上半部和下半部，使用这种划分作为边界以确保浅层和深层具有相同数量的特征。遵循[13, 24, 29]，作者采用了范数。与引入额外模块以减少冗余的方法[16, 11]相比，简单且无需成本。

目标函数。 总损失是标准交叉熵和骰子损失[23, 8, 1, 18]的组合，加权于和，在预测和 GT 分割结果之间进行评估。

其中和是平衡参数。

Dataset

突触多器官分割。遵循[1, 3, 8]，突触数据集包含30个案例，共3779张轴向腹部临床CT图像。其中，18个样本用于训练，而12个保留用于测试。作者使用平均Dice相似系数（DSC）作为评估指标，对八个腹部器官的性能进行评估。

自动化心脏诊断挑战数据集。 ACDC数据集包含来自不同患者的100个心脏MRI扫描，每个扫描都标记有左心室（LV）、右心室（RV）和心肌（MYO）。遵循[3, 17]，作者将70个案例（1930张轴向切片）用于训练，10个用于验证，20个用于测试。作者使用DSC作为评估指标来评估作者的方法。

图4：在Glas和MoNuSeg数据集上分割性能的比较。

图3：在Synapse数据集上分割性能的比较。

核分割和腺体分割。 腺体分割数据集（GlaS）[20]有85张图像用于训练和80张用于测试。多器官核分割（MoNuSeg）数据集[10]有30张图像用于训练和14张用于测试。遵循[23]，作者在GlaS和MoNuSeg数据集上执行三次5折交叉验证。平均DSC和IoU用作评估指标。### 实验设置

作者在这四个数据集上评估作者提出的损失在SwinUnet和Unet上的有效性，并且训练设置（即，批量大小、优化器、学习率等）与[1]一致。所有实验均使用224 x 224的输入图像大小，并采用与[8, 17, 23]相同的数据增强和预处理，使用带有24GB内存的Nvidia GTX3090进行训练。遵循[1, 18]，在SwinUnet中使用来自ImageNet的预训练权重，而Unet则从零开始训练。

与最先进方法的比较。 作者将作者的方法与最近的SOTA模型进行比较，包括R50 Unet [3]，Att-Unet [15]，Unet++ [30]，Transuet [3]，swinUnet [1]，Levit-Unet [26]，DeepLabv3 [4]，HiFormer [8]，PVT-cascade [17]，UCTransNet [23]，MedT [21]。

Results

表3：Glas和MoNuSeg数据集中不同方法的比较。

在Synapse数据集上的结果展示在表1中，ACDC数据集上的结果展示在表2中，Glas和MoNuSeg数据集上的结果展示在表3中。主要观察结果是，作者提出的损失函数是有效的，并能带来实质性的增益。特别是，通过利用作者提出的损失函数，Unet在四个数据集上的平均DSC分别提高了3.49%、1.75%、5.48%和3.73%。同样，在这四个数据集上，SwinUnet可以获得1.41%、3.42%、2.04%和1.69%的DSC提升。因此，作者增强的模型能够超越之前的现有技术水平（SOTA）方法。例如，作者增强的Unet和SwinUnet在平均DSC上分别展示了0.05%和0.25%的提升。与Hiformer、TransUnet、Unet和swinUnet [18, 1, 8, 3]相比，作者增强型模型的性能提升主要来自于解决具有挑战性的器官分割任务，如图3和表1所示，如胆囊、左肾、右肾和胰腺等。同样，在ACDC数据集中，与SOTA HiFormer [8]相比，作者的Unet和swinUnet方法在平均DSC上分别展示了0.6%和0.67%的提升。这种优势也推广到了Glas和MoNuSeg数据集，作者的方法在这两个数据集上相对于之前的SOTA方法UCTransNet分别实现了0.7%和1.1%的DSC提升。图4展示了在Glas数据集上作者方法的显著改进。需要注意的是，带有优化跳跃连接的Unet（例如，UCTTransNet [23]）导致了不相关的分割和不完整的形状。尤其是类似于背景的片段（图4的最后一行）。而带有作者提出的损失函数的SwinUnet展示了与真实情况非常接近的结果，具有完整的形状和清晰的背景，特别是在难以改进的样本上（图4的第三行）。这些观察结果表明，作者的方法可以提供上下文监督，确保跨块的正确语义并防止学习不相关的特征。

Ablation studies

为了深入了解作者的方法，作者使用Unet在Glas和MoNuSeg数据集上进行了系列消融研究。结果和分析如下：损失平衡超参数。如图5(a,b)所示结果，作者观察到和是最优设置。随着的权重减少（），性能迅速下降。这也表明，缺乏正确的语义监督是性能下降的原因之一。所提损失的有效性。作者进行了如图5(c,d)所示的采用不同损失的消融研究。两个数据集都显示出一致的结果，在使用这两种损失同时达到最优性能。