[论文解读] STA-Unet: Rethink the semantic redundant for Medical Imaging Segmentation

Vamsi Krishna Vasa, Wenhui Zhu+4 more|arXiv (Cornell University)|Oct 13, 2024

Radiomics and Machine Learning in Medical ImagingMedicine被引用 3

一句话总结

本文提出STA-Unet，一种通过引入超级标记注意力（Super Token Attention, STA）机制增强的U-Net架构，以减少视觉变换器浅层中的语义冗余，从而提升多器官分割性能。通过用稀疏、基于超像素的标记处理替代标准自注意力机制，STA-Unet在四个医学影像数据集上实现了最先进性能，Dice分数相比U-Net最高提升2.86%，相比TransUNet最高提升3.22%。

ABSTRACT

In recent years, significant progress has been made in the medical image analysis domain using convolutional neural networks (CNNs). In particular, deep neural networks based on a U-shaped architecture (UNet) with skip connections have been adopted for several medical imaging tasks, including organ segmentation. Despite their great success, CNNs are not good at learning global or semantic features. Especially ones that require human-like reasoning to understand the context. Many UNet architectures attempted to adjust with the introduction of Transformer-based self-attention mechanisms, and notable gains in performance have been noted. However, the transformers are inherently flawed with redundancy to learn at shallow layers, which often leads to an increase in the computation of attention from the nearby pixels offering limited information. The recently introduced Super Token Attention (STA) mechanism adapts the concept of superpixels from pixel space to token space, using super tokens as compact visual representations. This approach tackles the redundancy by learning efficient global representations in vision transformers, especially for the shallow layers. In this work, we introduce the STA module in the UNet architecture (STA-UNet), to limit redundancy without losing rich information. Experimental results on four publicly available datasets demonstrate the superiority of STA-UNet over existing state-of-the-art architectures in terms of Dice score and IOU for organ segmentation tasks. The code is available at \url{https://github.com/Retinal-Research/STA-UNet}.

研究动机与目标

• 解决基于变换器的U-Net架构浅层中固有的特征冗余问题，该问题会阻碍高效学习并增加计算开销。
• 探究在医学图像分割中应用视觉变换器时被忽视的语义冗余问题。
• 通过集成超级标记注意力（STA）机制提升U-Net性能，在最小化冗余注意力计算的同时保持丰富的语义信息。
• 在包括多器官分割和腺体/细胞核分割在内的多样化医学影像任务中，展示模型的泛化能力与鲁棒性。
• 提供一种计算效率更高的替代方案，用于过度参数化的基于变换器的模型，同时不牺牲分割精度。

提出的方法

• 将超级标记注意力（STA）机制集成到U-Net的编码器-解码器架构中，以替代标准多头自注意力层。
• 通过基于超像素的聚类将像素级特征转换为超标记，减少标记数量并聚焦于显著的视觉区域。
• 应用稀疏关联学习与标记空间映射，提升注意力效率并减少浅层中的计算量。
• 保留编码器与解码器之间的跳跃连接，但通过STA模块增强特征融合，以保留高分辨率上下文信息。
• 采用可学习的标记化策略，将局部图像块映射为紧凑且语义有意义的超标记，减少相邻相似像素之间的冗余注意力。
• 使用交叉熵损失与Dice损失端到端训练模型，以优化医学影像基准上的像素级分割性能。

实验结果

研究问题

• RQ1 基于变换器的U-Net模型浅层中的语义冗余在多大程度上损害了分割性能与效率？
• RQ2 超级标记注意力（STA）机制是否能有效减少早期层中的冗余注意力计算，同时保留关键的语义特征？
• RQ3 在多种医学影像数据集上，STA-Unet与最先进U-Net变体相比，在分割精度方面表现如何？
• RQ4 关键超参数（如标记大小和注意力头数量）对模型性能与效率的影响是什么？
• RQ5 所提出的架构是否能泛化到具有挑战性的分割任务中，如小尺寸或低对比度器官（如胰腺、肾脏）以及细粒度结构（如细胞核、腺体）？

主要发现

• 在ACDC数据集上，STA-Unet实现了92.25%的平均Dice分数，相比U-Net提升2.86个百分点，相比TransUNet提升2.83个百分点。
• 在MoNuSeg数据集上，STA-Unet取得81.06%的Dice分数，相比U-Net提升6.03个百分点，相比TransUNet提升3.22个百分点。
• 在Glas数据集上，STA-Unet实现91.03%的Dice分数，相比U-Net提升6.53个百分点，相比TransUNet提升2.97个百分点。
• STA-Unet显著提升了对小尺寸和具有挑战性器官（如胰腺、肾脏）的分割性能，优于SwinUNet等在这些结构上失败的模型。
• 视觉分析表明，STA-Unet在复杂区域（如腺体和细胞核）中产生的前景预测更准确、更完整，优于SwinUNet和LeViT-UNet。
• 消融实验表明，标记大小与注意力头数量均显著影响模型性能，最优设置在精度与计算成本之间实现良好平衡。