[论文解读] Multi-dimension Transformer with Attention-based Filtering for Medical Image Segmentation

Wentao Wang, Xi Xiao+6 more|arXiv (Cornell University)|May 20, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

本文提出 MDT-AF，一种基于注意力过滤的多维变换器，用于医学图像分割，通过在补丁嵌入阶段整合基于注意力的噪声过滤，并将自注意力机制扩展至空间和通道维度，从而增强特征表示。该方法在三个公开基准上实现了最先进性能，提升了低信噪比条件下的鲁棒性与更丰富的特征学习能力。

ABSTRACT

The accurate segmentation of medical images is crucial for diagnosing and treating diseases. Recent studies demonstrate that vision transformer-based methods have significantly improved performance in medical image segmentation, primarily due to their superior ability to establish global relationships among features and adaptability to various inputs. However, these methods struggle with the low signal-to-noise ratio inherent to medical images. Additionally, the effective utilization of channel and spatial information, which are essential for medical image segmentation, is limited by the representation capacity of self-attention. To address these challenges, we propose a multi-dimension transformer with attention-based filtering (MDT-AF), which redesigns the patch embedding and self-attention mechanism for medical image segmentation. MDT-AF incorporates an attention-based feature filtering mechanism into the patch embedding blocks and employs a coarse-to-fine process to mitigate the impact of low signal-to-noise ratio. To better capture complex structures in medical images, MDT-AF extends the self-attention mechanism to incorporate spatial and channel dimensions, enriching feature representation. Moreover, we introduce an interaction mechanism to improve the feature aggregation between spatial and channel dimensions. Experimental results on three public medical image segmentation benchmarks show that MDT-AF achieves state-of-the-art (SOTA) performance.

研究动机与目标

• 解决医学图像中低信噪比带来的挑战，该问题会阻碍特征学习与分割精度。
• 提升视觉变换器在捕捉复杂空间与通道依赖关系方面的能力，以支持医学图像分割。
• 克服标准自注意力机制在建模医学图像中长距离与多维特征时的局限性。
• 开发一种鲁棒的端到端框架，通过提升特征质量与分割精度，而无需依赖高级训练策略。

提出的方法

• 在补丁嵌入阶段引入基于注意力的过滤机制，以粗到精的方式精炼粗粒度特征并抑制噪声。
• 重新设计补丁嵌入模块，生成引导注意力权重，以在自注意力计算前过滤无关或噪声特征。
• 将自注意力机制扩展至空间、通道及联合空间-通道维度，实现更丰富的特征表示。
• 实现一个多维变换器模块，用于在空间与通道特征之间进行特征交互与聚合，以增强判别能力。
• 采用多级MLP解码器，融合编码器中的层次化特征，以实现精确的语义分割掩码预测。
• 将基于注意力的过滤与多维自注意力机制整合至统一的编码器-解码器架构中，支持端到端训练。

实验结果

研究问题

• RQ1 在补丁嵌入阶段引入基于注意力的过滤是否能提升低信噪比医学图像中特征的质量与抗噪鲁棒性？
• RQ2 将自注意力机制扩展至空间与通道维度是否能增强医学图像分割中的特征表示与分割性能？
• RQ3 与现有视觉变换器及基于CNN的方法相比，所提出的MDT-AF框架在多样化医学图像基准上的有效性如何？
• RQ4 多维自注意力机制在建模医学图像中复杂解剖结构方面的提升程度如何？

主要发现

• 在Lung X-ray数据集上，MDT-AF的Dice相似系数（DSC）达到97.17%，较基线提升+0.33%。
• 在Skin Lesion数据集上，MDT-AF的DSC达到94.25%，较基线提升+0.09%。
• 在Kvasir-SEG数据集上，MDT-AF的DSC达到93.38%，较基线方法提升+1.11%。
• 消融实验表明，基于注意力的过滤与多维自注意力组件均对性能提升有显著贡献。
• 完整MDT-AF组件的模型在Kvasir-SEG数据集上，相比无过滤的变体（DSC提升+1.11%）与无多维注意力的变体（DSC提升+0.78%）均表现更优。
• 可视化对比显示，MDT-AF在噪声较多的内窥镜图像中能生成更清晰的边界与更少的误分类，优于其他方法。