STA-Unet 医学影像分割语义冗余的再思考

近年来，在医学影像分析领域，利用卷积神经网络（CNNs）取得了显著的进步。特别是基于U型结构的深度神经网络（UNet）及其 Shortcut 已在多种医学影像任务中得到应用，包括器官分割。 尽管它们取得了巨大的成功，但CNNs在学习全局或语义特征方面表现不佳。特别是需要人类般推理才能理解上下文的特征。许多UNet架构试图通过引入基于Transformer的自注意力机制进行调整，并取得了显著的性能提升。 然而，Transformer天生存在冗余，在浅层学习时需要冗余地学习，这通常会导致附近像素的注意力计算增加，这些像素提供的信息有限。 最近引入的超级像素注意力（STA）机制将超级像素的概念从像素空间转换到 Token 空间，使用超级像素作为紧凑的视觉表示。这种方法通过在视觉 Transformer 中学习有效的全局表示，特别是在浅层，来解决冗余问题。 在本研究中，作者在UNet架构中引入了STA模块（STA-UNet），在限制冗余的同时保持丰富的信息。在四个公开可用的数据集上的实验结果表明，与现有最先进的架构相比，STA-UNet在器官分割任务上的Dice分数和IOU方面具有优势。 代码：https://github.com/Retinal-Research/STA-UNet。

1 Introduction

利用深度学习技术的进步，计算机视觉技术已成为医学图像分析的重要组成部分。在这些技术中，图像分割具有重大意义。具体而言，精确可靠的医学图像分割是计算机辅助诊断和影像引导手术操作的基础组成部分[5, 10]。

当前医学图像分割的主要方法是利用具有U型结构的完全卷积神经网络（FCNNs）。广为人知的U-Net[30]是这种结构的典型示例，它采用对称的编码器-解码器设计，通过 Shortcut 相互连接。编码器通过多个卷积和降采样层提取具有广泛感受野的深度特征。解码器则将这些深度特征上采样回原始分辨率，以实现精确的像素级语义预测，同时 Shortcut 将编码器中不同尺度的特征融合，以减轻下采样导致的空间信息损失。这种精心设计的设计使得U-Net在多个医学影像任务上取得了显著的成功。这些基于FCNN的方法在心脏分割、器官划分和病变检测方面的出色性能强调了卷积神经网络在学习区分特征方面的强大能力。

尽管基于CNN的技术在医学图像分割方面取得了令人瞩目的成果，但它们仍然无法达到临床应用所需的高精度标准。医学图像分割仍然是一个具有挑战性的问题，主要原因在于卷积操作固有的局部特征关注特性，这使得CNN难以捕捉明确的全球和长程语义交互。

最近，受到自然语言处理（NLP）[34]中Transformer的成功启发，研究行人开始探索其在视觉领域[4,8,27]的应用，以解决使用自注意力机制的CNN的局限性。尽管Vision Transformer（ViT）在大规模感受野下能够捕捉图像 Patch 之间的长程依赖性，但其由于缺乏固有的局部性而面临着保留细粒度局部上下文方面的挑战。

为解决此问题，最近的方法 提出了将CNN和ViTs相结合的混合模型，这些模型显著增加了计算复杂性和参数数量。深度学习中的超参数化是一个普遍问题，经常导致特征冗余和次优特征表示 。

尽管其影响重大，但现有研究并未充分探索或考虑这个问题。除了上述方法外，一些方法旨在提高UNet的架构设计。例如，Att-UNet [28]引入了基于注意力的 Shortcut 来过滤无关特征，而UNet++ [39]用嵌套密集跳跃路径替换了传统的 Shortcut （如连接或加法）。UCTransNet [35]对各种 Shortcut 策略进行了深入分析，并提出了使用通道 Transformer 作为传统 Shortcut 的替代方案。最近提出的Seg-Swinunet [40]利用具有最高语义内容（即解码器的最后一层）的特征图向其他块提供额外的监督，通过特征蒸馏减少特征冗余。然而，作者从不同的角度研究这个问题。

作者的初步分析表明，Transformer UNet架构的浅层块之间存在显著的相似性。这一观察意味着模型在浅层表现出一种惯性学习模式，导致无法有效地捕捉和编码复杂上下文信息。现有研究很少关注这一固有局限性。[15]将像素域的超像素概念转移到 Token 域，将超 Token 视为视觉信息的简洁表示。这种方法将稀疏关联学习、自注意力和 Token 空间映射集成在一起，以提高视觉 Token 处理效率，从而实现丰富的特征学习。在本研究中，作者试图通过将超 Token 注意力集成到UNet架构中并提高多器官分割挑战的性能来解决冗余问题。

作者的工作主要有三点贡献：

(i) 作者强调了基于 Transformer UNet的浅层冗余性，以促进该领域的研究。

(ii) 作者将超级 Token 注意力（STA）块集成到UNet架构中，以最小化其他基于 Transformer UNet模型中观察到的冗余性，同时保留对有效学习所需丰富的语义信息。

(iii) 作者在四个公开可用的医学影像数据集上进行了全面的评估，结果表明提出的方法在器官分割任务上优于现有的相关最先进方法。

2 Related Work

基于UNet的架构： 医学图像分割的早期方法主要依赖于轮廓为基础的方法和传统的机器学习技术[12, 33]。然而，深度卷积神经网络（CNNs）的出现带来了显著的进步，其中UNet[30]专门针对医学图像分割而设计。U-Net独特的U形架构，因其简单和卓越的性能而闻名，已经激发了多种变体，包括ResNet[37]，Dense-Unet[25]，U-Net++[39]，和UNet3+[14]。基于CNN的架构在捕捉冗余信息方面存在缺陷，并且没有关注画布上不同区域之间的依赖关系。

基于Transformer的UNet架构： Vaswani等人[34]在自然语言处理中引入了自注意力机制，用于衡量不同单词之间的相对重要性。这一进展促成了视觉 Transformer （ViT）[8, 27]的出现，该 Transformer 将Transformer架构应用于图像处理任务，以实现与CNN相当的成功。这些 Transformer 与UNet设计[3, 5, 38]相结合，旨在结合CNN和Transformer的优势。

陈等人[5]尝试将UNet架构中的Transformer结合在一起。编码器中的Transformer将CNN特征图中的图像块分词，以捕捉全局上下文。同时，解码器将编码特征上采样并将其与CNN的高分辨率特征图合并。尽管ViT由于其大的感受野在捕捉图像块（分词）之间的长期依赖关系方面表现出色，但由于其缺乏固有的局部性，因此在保持详细局部上下文方面面临挑战。为了克服这一限制，Swin-Unet[3]使用移动窗口分词[27]调整了注意力机制。这使得模型可以将窗口注意力限制在局部区域。尽管这种适应限制了冗余性，但它在浅层中并未完全消除。同时，朱等人[40]提出了Seg-SwinUNet，通过在编码器和解码器之间平衡监督和减少特征冗余，解决了UNet在医学图像分割方面的性能问题。该方法通过使用特征蒸馏从最语义丰富的特征图提供额外的监督，以最小化计算开销，提高分割精度。然而，该工作仍然局限于Swin-Unet，尚未进行进一步的研究，以将这种方法与其他架构相结合。

许等人[38]提出了LeViT-Unet，其中LeViT[9]作为LeViT-UNet框架内的编码器，因为它在Transformer块中有效地平衡了准确性和效率。此外， Shortcut 将来自Transformer和LeViT卷积块的多尺度特征图集成到解码器中。

由于LeViT在保留和传递信息到解码器方面起核心作用，因此无法避免冗余 Token 信息，导致计算成本增加。最近提出的Hiformer[11]将CNN和Transformer架构集成起来，用于捕捉医学图像分割的局部和全局特征。它使用Swin Transformer和基于CNN的编码器，通过在编码器-解码器结构中的双层融合（DLF）模块将多尺度特征表示结合在一起。大量实验表明，与其它方法相比，HiFormer在准确性和效率方面具有优越性能。

减轻特征冗余： Oktay等人[28]提出了注意门（AG）来通过抑制无关区域和突出重要特征，关注各种形状和大小的目标结构。这样就无需使用外部的局部化模块，因为AG可以很容易地集成到CNN架构（如U-Net）中，且计算成本最小。周等人[39]提出了UNet++架构，用于深度监督连接编码器和解码器的编码器-解码器网络，通过嵌套、密集的跳过路径。这些重新设计的路径旨在减少编码器特征图和解码器特征图之间的语义差距，使优化器的学习任务更容易。

王等人[35]提出了U-TransNet，用通道 Transformer （CTrans）模块替代传统U-Net的 Shortcut ，该模块包括两个子模块：通道与 Transformer （CCT）用于多尺度通道融合，以及通道内交叉注意力（CCA）将融合的特征引导到解码器。这种新的连接结构解决了编码器和解码器特征之间的语义差距，以提高分割效果。朱等人[40]提出了一种在编码器和解码器之间平衡监督，并在UNet中通过解码器最语义丰富的特征图（解码器最后一层）提供额外监督来减少特征冗余的方法，并利用特征蒸馏最小化冗余信息并提高学习效率。这种方法可以无缝集成到现有的UNet架构中，且计算成本最小，从而在各种医学图像分割任务上提高性能。

3 Preliminary Analysis

在这里，作者采用中心核对齐（CKA）[19]来研究最近流行的U-net模型，包括SwinUnet、LeViT-Unet、TransUnet和HiFormer。这种技术使作者能够在层大小不同时计算块级相似性。块级相似性矩阵能够提供关于不同神经网络架构在训练过程中如何学习和表示信息（或不同块）的见解。

在数学上，给定两组表示 和 ，作者首先通过径向基函数（RBF）计算它们的核矩阵 和 。

其中 和 分别表示矩阵的第 行和第 列元素。 分别表示集合 的第 个和第 个样本。

随后，通过 RBF-CKA 计算相似度矩阵，如下所示：

其中 表示校正矩阵。 是一个具有 形状的单位矩阵，其中 表示集合中的样本数量。

在图1中，作者呈现了作者的调查结果，其中非对角元素中的高值通常表示对应块之间的强烈相似性。在各种架构中，相似度矩阵通常在浅层块之间显示出较高的相似性，这表明这些早期层具有较高的冗余。这立即暗示神经网络在浅层块中懒惰，无法学习丰富的信息。这一观察使作者努力解决并减少这种冗余，这在所提出的工作中得到了体现。

4 Method

Super Token Attention Module

根据第3节中提出的分析，基于transformer架构的浅层存在明显的冗余，这导致了信息保留的低效。通过自监督学习，超级 Token [15]可以缓解这一缺陷，通过学习有效的全局表示。超级 Token 被视为是视觉信息的简洁表示，通过从像素域到 Token 域调整超像素[17]的概念。这种方法结合稀疏关联学习、自注意力以及 Token 空间映射来提高视觉 Token 处理的效率。作者在UNet架构中重新引入了超级 Token 注意力（STA模块），以利用其在医学图像分割任务中的优势。

Token 和超级 Token 之间的关联旨在计算从到的关联()并更新超级 Token 。作者使用公式5计算以获得注意力类型的权重。

作者还在图4中可视化了从现有的基于 Transformer 的UNet架构中获得的注意力图。值得一提的是，超级 Token 注意力在浅层块中对较小的感兴趣区域（如肾脏和主动脉）赋予了更高的权重。

STA-UNet architecture

作者将本节奉献给简要概述所提出的Unet架构（如图3所示）。与任何其他UNet架构一样，所提出的模型包括编码器、解码器、 Bottleneck 和 Shortcut 。在这个架构中的关键性能增强器是集成在每个编码器和解码器阶段的超级 Token 注意力（STA）模块。相反，利用在最新UNet架构[3, 5, 38]中使用的Transformer模块，作者在卷积层中实现维数变化，在应用注意力机制之前过滤关键信息。

输入图像被下采样到一半尺寸()，并在每个阶段将通道（C）维数加倍。在STA模块（在CPE阶段）中提取位置嵌入，然后生成超级 Token ，并计算 Token 与超级 Token 之间的相关性，如第4.1节所述。对称解码器架构通过结合上采样块（以恢复原始图像形状）和STA模块进行调整。

在处理过程中提取的上下文特征通过 Shortcut 与编码器的多阶段特征进行拼接。这种融合可以减轻下采样通常会导致的空间信息损失，从而增强模型保留细粒度细节的能力。每个组件的详细信息在本节中记录。

编码器 编码器每个阶段有两个组成部分。降采样块后紧跟着一个STA模块。降采样块处理正向传播过程中的维数降低，因为作者不依赖[3]中的卷积块合并阶段。降采样块由2层卷积和批量归一化组成。卷积层具有3x3的核大小和步长为2，内插为1。然后作者对输出特征使用ReLU [1]激活并进行降维，使用最大池化层。为了保留这个阶段的完整空间信息，作者将ReLU激活的输出通过 Shortcut 传递给解码器，而不是传统的Unet架构。

解码器借鉴了[30]中的设计思想，解码器设计为对称于编码器。解码器也由两个组件组成，即上采样块和STA模块。上采样块包括一个卷积层，类似于编码器（第4.2节）中讨论的卷积层，用于增加输入特征的空间维度。

作者将通过 Shortcut 从编码器获取的特征图与来自前一个解码阶段（STA模块的输出）的特征图进行拼接。编码器（在池化之前）的完整空间信息有助于作者将空间信息与注意力机制的上下文特征相结合，从而提高学习效果。

然后，作者将结果输入上采样块。解码器将输出传递到输出投影层，以获得输入图像维数和与分割类数相等的通道大小。输出通过C轴的Softmax层处理，以获得类概率。

Bottleneck 由2个卷积层和批量归一化（BatchNorm）层组成，输出使用ReLU激活函数。

5 Experiments and results

Datasets

作者在四个公开可用的数据集上验证了所提出方法的有效性：Synapse多器官分割（Synapse Multi-Organ Segmentation），自动心脏诊断挑战（ACDC）数据集[2]，核分割（MoNuSeg）[20, 21]和结肠组织切片图像中的腺体分割（GlaS）[31]。遵循[3, 5, 11]的研究，作者在Synapse多器官分割数据集上训练了所提出的模型。该数据集包括30个案例，总共包含3,779个横断面腹部CT图像。作者为13个腹部器官提供了分割 Mask ，其中作者用9个类别训练所提出的模型。在模型开发方面，作者分配了18个案例用于训练，而12个案例用于测试。性能基于分割的8个腹部器官，使用平均Dice相似系数（DSC）作为主要评估指标。

ACDC数据集包括来自不同患者群体的100个心脏MRI扫描，并带有左心室（LV）、右心室（RV）和心肌（Myo）的标注。根据先前的研究[29, 5]，作者将数据集划分为70个案例（1,930个横断面切片）用于训练，10个案例用于验证，20个案例用于测试。

作者使用Dice相似系数（DSC）作为评估指标来评估所提出方法的表现。GlaS [32]和MoNuSeg [22]数据集是显微镜图像的集合。GlaS数据集包含85个用于训练的图像和80个用于测试的图像。MoNuSeg数据集包括30个用于训练的图像和14个用于测试的图像。作者使用平均Dice相似系数（DSC）和交点与 Union（IoU）作为指标来评估后两个数据集的表现。

Implementation details

作者遵循了简单易复现的训练方案 [3, 5, 36, 40]。Synapse CT 数据集包含 3D CT 扫描图像，每个切片在灰度域进行映射。为了在这个数据集上进行训练，作者提取每个切片并中心裁剪，以保留输入的 224 x 224 图像。

作者使用随机梯度下降（SGD）优化器训练模型，以实现更平滑的收敛。批量大小设置为 8。初始学习率 设置为 。

每个 epoch 的迭代学习率 由公式 9 确定，其中 t 表示当前迭代次数，N 表示一个 epoch 中的最大迭代次数。

作者训练模型收敛于交叉熵和Dice损失的总和，分别保持0.4和0.6的权重。为解决数据集有限的问题，作者采用了以下数据增强特性：随机翻转（水平）和旋转，概率为0.5。作者遵循了ACDC数据集的相同实验设置。对于GlaS和MoNuSeg数据集，批处理大小为18。

作者使用的初始学习率，并使用余弦调度器更新学习率。作者在Nvidia RTX3090的计算能力上进行实验，具有24G内存。

作者将作者的方法与最近的最先进模型进行比较，包括UNet [30]，R50 U-Net [7]，Att-UNet [28]，TransUNet [5]，SwiUNet [3]，LeViT-UNet [38]，HiFormer [11]，以及SegSwinUNet [40]。

Ablation Study

理解单个参数对模型性能的影响对于确定最优架构至关重要。为了深入了解作者提出的模型中变大小的 Token 和自注意力头的影响，作者在GlaS数据集上进行了消融研究。

在基于Transformer架构的模型中，增加注意力头通常会使模型从更多的区域捕捉信息，并基于决策确定其重要性。从表2可以看出，超 Token 注意力机制可以在有限的注意力头下实现优越的性能，降低关注或冗余的交集。这在图5(a)中也有所体现。作者选择根据作者提出的方法的性能来突出显示的注意力头趋势。

Token 大小会影响模型捕捉空间细节和上下文信息的能力。较大的 Token 提供更广泛的上下文，但可能会降低分辨率，而较小的 Token 可以捕捉更精细的细节，但可能会增加计算复杂性。这一点在表3中同样明显。平衡 Token 大小对于优化模型性能和效率至关重要。

但在提出的超级 Token 注意力方法中，作者发现性能随着 Token 大小的变化非常微小，因此限制了性能对 Token 大小的依赖性；这一点在图5（b）中同样明显。基于研究，作者选择了每秒浮点运算次数相对较低的模型，以降低计算复杂性。

Results

表4报告了Synapse数据集的性能分析，表5报告了ACDC数据集的性能分析，表6报告了Glas和MoNuSeg数据集的性能分析。作者的主要结论是，作者提出的架构是有效且计算合理的，并在定量指标上取得了显著改进。作者将相对于UNet[30]和首次提出的基于Transformer的UNet架构TransUNet[5]的增益/损失报告为百分比。

在四个数据集（相对于UNet）中，作者的方法分别实现了4.99%，2.86%，6.53%和6.03%的改进。同样，与TransUNet相比，作者的方法分别实现了4.14%，2.83%，2.97%和3.22%的改进。与最近建立的工作如HiFormer[11]和Seg-Swinunet[40]相比，作者在DSC方面分别实现了0.49%和0.18%的改进。这种巨大的DSC改进源于更准确地分割了肾脏（L&R）和胰腺等困难器官。

如图6所示，作者关注到的胰腺和胃的分割（图6的黄色部分，第一行）。值得注意的是，SwinUNet无法分割它们，其他模型如[5, 38, 40]也没有完全分割胰腺。所提出的模型很好地解决了这个挑战，与Hi-Former相当。

作者通过提高DSC在ACDC、Glas和MoNuSeg数据集上的2.86%、6.53%和6.03%，分别比UNet提高了2.86%、6.53%和6.03%（与UNet相比），来验证作者工作的泛化性。作者在ACDC和MoNuSeg数据集的分割任务上超越了所有基于Transformer的UNet架构，并且在Glas数据集上以0.64%的微小差异排名第二，仅次于UNet（见表5和6）。作者在图7中通过使用Glas（使用Glas）和Nuclear（使用MoNuSeg）对分割性能进行了视觉比较。

从SwinUNet和LeViT-UNet（图7中的顶部两行）的预测中可以看出，Gland分割的前景分类问题明显，这个问题被提出的超 Token 注意力（STA-UNet）显著解决，从而实现了精确的分割。作者还突出了在Glas数据集中区分前景和背景的困难。在MoNuSeg的情况下，作者提出的模型实现了与 GT 值高度可比的结果，捕捉到了完整的形状并保持了清晰的背景，即使在具有挑战性的样本中也是如此（如图7中的第三行所示）。

这些发现进一步加强了作者的主张，即STA-UNet可以提高分割性能，即使是在Transformer基础架构中通常出现的浅层特征冗余减少的情况下。

6 Conclusion

在这项研究中，作者重新引入了在UNet架构中的超级令牌注意力（STA）作为一个STA模块，以解决现有基于变换器架构中固有的特征冗余问题，同时在器官分割任务上提升性能。

作者报告了初步分析，以数学方式使基于变换器架构中的冗余问题变得明显，并鼓励研究以减轻这一问题。作者的研究发现，在四个公开可用的数据集上，STA-UNet相较于现有基准取得了显著的改进，证明了STA-UNet在医学图像分割中的潜力。

作者进行的广泛消融研究解释了两个主要参数变化，即令牌大小和注意力头数量，对性能的影响。虽然作者的实验仅限于多器官分割任务，但STA-UNet具有更广泛应用的潜力，例如在各种医学数据集中的异常检测和恢复。

作者预计在未来的研究中将进一步探索所提架构在其他医学图像应用中的效用。

参考文献

[0]. STA-Unet: Rethink the semantic redundant for Medical Imaging Segmentation.