[论文解读] U-Net Transformer: Self and Cross Attention for Medical Image Segmentation

Olivier Petit, Nicolas Thome+2 more|arXiv (Cornell University)|Mar 10, 2021

Radiomics and Machine Learning in Medical Imaging被引用 18

一句话总结

本文提出U-Transformer，一种基于U-Net的架构，通过引入多头自注意力和交叉注意力机制，提升医学图像分割性能，尤其适用于低对比度及复杂器官。通过建模长距离上下文依赖关系，并利用注意力引导的跳跃连接优化空间细节，U-Transformer在腹部CT数据集上实现最先进性能，相较于U-Net和注意力U-Net，在胰腺等挑战性器官上的Dice分数最高提升3.4%。

ABSTRACT

Medical image segmentation remains particularly challenging for complex and low-contrast anatomical structures. In this paper, we introduce the U-Transformer network, which combines a U-shaped architecture for image segmentation with self- and cross-attention from Transformers. U-Transformer overcomes the inability of U-Nets to model long-range contextual interactions and spatial dependencies, which are arguably crucial for accurate segmentation in challenging contexts. To this end, attention mechanisms are incorporated at two main levels: a self-attention module leverages global interactions between encoder features, while cross-attention in the skip connections allows a fine spatial recovery in the U-Net decoder by filtering out non-semantic features. Experiments on two abdominal CT-image datasets show the large performance gain brought out by U-Transformer compared to U-Net and local Attention U-Nets. We also highlight the importance of using both self- and cross-attention, and the nice interpretability features brought out by U-Transformer.

研究动机与目标

• 解决U-Net在捕捉低对比度、复杂解剖结构中长距离上下文依赖关系和空间关系方面的局限性。
• 提升对小尺寸及视觉上模糊器官（如胰腺、胆囊和胃）的分割精度。
• 将自注意力与交叉注意力机制整合到U-Net框架中，以建模全局上下文并优化空间细节。
• 通过注意力图可视化与消融研究，提升模型可解释性。
• 验证自注意力与交叉注意力在统一的Transformer增强U-Net架构中所展现的互补优势。

提出的方法

• 提出一种U-Transformer网络，将U-Net编码器-解码器架构与编码器末端的多头自注意力（MHSA）结合，以建模整个图像的全局上下文交互。
• 在跳跃连接中引入多头交叉注意力（MHCA），融合高层语义特征与高分辨率特征，实现细粒度的空间细节恢复。
• 使用可学习的位置编码，保留MHSA与MHCA模块中的空间位置信息，提升注意力定位能力。
• 通过在解码器不同阶段堆叠MHCA模块，实现多层级交叉注意力，逐步优化预测结果。
• 采用双分支注意力机制：MHSA捕捉特征间的长距离依赖，MHCA在特征上采样过程中过滤非语义噪声。
• 采用标准U-Net主干网络，结合残差块与跳跃连接，通过引入注意力模块而非替换卷积层来增强网络性能。

实验结果

研究问题

• RQ1 编码器中的自注意力是否能通过建模超出U-Net有限感受野的长距离上下文依赖关系，提升复杂器官的分割性能？
• RQ2 跳跃连接中的交叉注意力是否能通过过滤非语义特征并整合位置信息，增强空间细节恢复能力？
• RQ3 自注意力与交叉注意力机制在单独使用与联合使用时，其分割性能与可解释性表现如何比较？
• RQ4 位置编码与多层级注意力对解码路径中交叉注意力性能的影响如何？
• RQ5 自注意力与交叉注意力机制的联合集成是否能提升模型在各类器官上的泛化能力，特别是低对比度与复杂形状器官？

主要发现

• 在IMO数据集的胰腺上，U-Transformer的Dice相似系数（DSC）达到73.10%，相较于U-Net提升+3.4%，相较于注意力U-Net提升+4.45%。
• 在TCIA胰腺数据集上，U-Transformer的DSC达到80.65%，显著优于U-Net（76.35%）与注意力U-Net（77.23%）。
• U-Transformer中自注意力与交叉注意力的联合使用优于单一机制，证明二者具有互补性。
• 位置编码使TCIA数据集上MHCA性能提升+1.7分，IMO数据集上提升+0.6分，凸显其在空间定位中的重要性。
• 多层级交叉注意力使TCIA数据集性能提升+1.8分，IMO数据集提升+0.6分，表明分层注意力优化可有效提升分割精度。
• Hausdorff距离结果与Dice分数趋势一致，表明U-Transformer能减少预测伪影并提升边界精确度。