arXiv:2410.11578eess.IVcs.AI2024-10被引 7

用超令牌机制减少医学图像分割中的冗余计算,提升模型效率与精度。

STA-Unet: Rethink the semantic redundant for Medical Imaging Segmentation

  • 在UNet中引入超令牌注意力模块,压缩浅层特征冗余。
  • 在4个公开数据集上,Dice分数和交并比均优于现有方法。
  • 适合需要高效高精度分割的医学影像研究者使用。

近年来,基于卷积神经网络(CNN)的医学图像分析取得了显著进展。尤其是采用跳跃连接的U形架构(UNet)被广泛应用于器官分割等任务。尽管表现优异,但CNN难以学习全局或语义特征,尤其在需要类人推理理解上下文的任务中表现不足。为改进此问题,许多UNet架构引入了基于Transformer的自注意力机制,性能得到提升。然而,Transformer在浅层存在固有的冗余问题,导致对邻近像素的注意力计算量大但信息增益有限。最近提出的超级令牌注意力(STA)机制将超像素概念从像素空间映射到令牌空间,以超令牌作为紧凑的视觉表示,有效缓解了这一冗余问题,提升了视觉变换器中全局表征的学习效率。本文在UNet架构中引入STA模块(STA-UNet),在不丢失丰富信息的前提下限制冗余。在四个公开数据集上的实验表明,STA-UNet在器官分割任务中,于Dice分数和交并比方面均优于现有先进模型。代码已开源: https://github.com/Retinal-Research/STA-UNet。

原文摘要 · Abstract (English)

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}.

医学图像注意力机制分割

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。