[论文解读] Efficient Folded Attention for 3D Medical Image Reconstruction and Segmentation

Hang Zhang, Jinwei Zhang+5 more|arXiv (Cornell University)|Sep 13, 2020

Medical Image Segmentation Techniques参考文献 34被引用 8

一句话总结

本文提出了一种用于3D医学图像重建与分割的新型折叠注意力（FA）机制，显著降低了计算成本和GPU显存使用量，同时保持或提升了准确率。通过使用四种排列方式对张量进行折叠与展开，以近似全连接关联矩阵，FA实现了空间与通道依赖关系的高效联合建模，在定量磁化率成像（QSM）和多发性硬化症病灶分割任务中达到最先进性能，相比标准注意力机制，显存使用量降低达97.9%，FLOPs降低达88.9%。

ABSTRACT

Recently, 3D medical image reconstruction (MIR) and segmentation (MIS) based on deep neural networks have been developed with promising results, and attention mechanism has been further designed to capture global contextual information for performance enhancement. However, the large size of 3D volume images poses a great computational challenge to traditional attention methods. In this paper, we propose a folded attention (FA) approach to improve the computational efficiency of traditional attention methods on 3D medical images. The main idea is that we apply tensor folding and unfolding operations with four permutations to build four small sub-affinity matrices to approximate the original affinity matrix. Through four consecutive sub-attention modules of FA, each element in the feature tensor can aggregate spatial-channel information from all other elements. Compared to traditional attention methods, with moderate improvement of accuracy, FA can substantially reduce the computational complexity and GPU memory consumption. We demonstrate the superiority of our method on two challenging tasks for 3D MIR and MIS, which are quantitative susceptibility mapping and multiple sclerosis lesion segmentation.

研究动机与目标

• 解决传统自注意力机制在3D医学图像处理中计算成本高、显存占用大的问题。
• 开发一种高效注意力机制，能够在不增加计算复杂度的前提下建模空间与通道依赖关系。
• 降低基于3D U-Net的网络中注意力模块在医学图像任务中的显存与FLOPs开销。
• 实现注意力机制在大规模3D医学图像任务（如QSM重建与MS病灶分割）中的实际部署。

提出的方法

• 该方法引入张量折叠与展开操作，将3D特征张量转换为更小的子张量，从而计算四个子关联矩阵，而非一个大型关联矩阵。
• 通过特征张量的四种排列方式生成子关联矩阵，以较低复杂度近似全连接全局注意力矩阵。
• 每个子关联矩阵通过查询-键点积运算后接Softmax计算，结果通过可学习投影进行聚合。
• 折叠注意力模块通过将每个特征元素视为所有通道上的向量，实现空间与通道依赖关系的联合建模。
• 该方法对关联矩阵施加秩一约束，以增强泛化能力并减少过拟合。
• FA模块为即插即用设计，可无缝集成至现有3D CNN架构（如U-Net），仅需极少架构修改。

实验结果

研究问题

• RQ1 能否设计一种计算高效的注意力机制，以建模3D医学图像中的空间与通道依赖关系？
• RQ2 张量折叠与展开在多大程度上可降低3D医学图像任务中自注意力机制的FLOPs与GPU显存使用量？
• RQ3 所提出的折叠注意力机制是否能在具有挑战性的3D医学图像重建与分割基准上保持或提升性能？
• RQ4 在临床场景中缺乏真实标签的情况下，折叠注意力模块对未见数据的泛化能力如何？

主要发现

• 在定量磁化率成像（QSM）任务中，FA-Net的RMSE最低，仅为31.18，优于QSMnet、DA-Net与RSA-Net。
• 在QSM任务中，FA-Net的PSNR最高，达到49.06，表明重建保真度更高，整体误差更小。
• 在QSM任务中，FA-Net的HFEN最低，为32.49，表明对高频细节与结构相似性的保留更优。
• 在QSM任务中，FA-Net的SSIM最高，达到0.9833，表明与真实值在结构与对比度相似性方面表现更优。
• 与标准自注意力（SA）相比，FA模块将GPU显存使用量降低97.9%，FLOPs降低88.9%，且准确率相当或更优。
• 定性结果表明，即使在缺乏真实COSMOS数据的情况下，FA-Net在MS患者扫描中仍能生成更清晰的病灶边界与更多高信号病灶。