[论文解读] Cross-modal Attention for MRI and Ultrasound Volume Registration

Xinrui Song, Hengtao Guo+7 more|arXiv (Cornell University)|Jul 9, 2021

Advanced Neural Network Applications参考文献 19被引用 6

一句话总结

本文提出了一种跨模态注意力机制，显式建模磁共振成像（MRI）与超声波图像体积之间的空间对应关系，以提升医学图像配准性能。通过将该注意力模块集成到轻量级卷积神经网络中，该方法在仅使用之前沿SOTA模型1/10参数量和一半推理时间的前提下，实现了最先进水平的配准精度——将配准误差从10.17±5.75 mm降低至3.71±1.99 mm。

ABSTRACT

Prostate cancer biopsy benefits from accurate fusion of transrectal ultrasound (TRUS) and magnetic resonance (MR) images. In the past few years, convolutional neural networks (CNNs) have been proved powerful in extracting image features crucial for image registration. However, challenging applications and recent advances in computer vision suggest that CNNs are quite limited in its ability to understand spatial correspondence between features, a task in which the self-attention mechanism excels. This paper aims to develop a self-attention mechanism specifically for cross-modal image registration. Our proposed cross-modal attention block effectively maps each of the features in one volume to all features in the corresponding volume. Our experimental results demonstrate that a CNN network designed with the cross-modal attention block embedded outperforms an advanced CNN network 10 times of its size. We also incorporated visualization techniques to improve the interpretability of our network. The source code of our work is available at https://github.com/DIAL-RPI/Attention-Reg .

研究动机与目标

• 解决在前列腺癌穿刺活检中，MRI与经直肠超声（TRUS）图像体积之间准确跨模态配准的挑战。
• 克服传统卷积神经网络在建模不同成像模态之间特征空间对应关系方面的局限性。
• 提出一种新型注意力机制，显式捕捉MRI与TRUS图像体积之间全局特征的对应关系。
• 通过显著更小、更高效的网络结构，相比现有深度学习方法，提升配准性能。
• 利用Grad-CAM等可视化技术，增强配准网络的可解释性。

提出的方法

• 提出一种跨模态注意力模块，将非局部注意力机制扩展至在MRI（固定）与TRUS（移动）图像体积的特征之间运行。
• 该注意力模块计算一个体积中每个特征与另一体积中所有特征之间的相似性，实现全局特征匹配。
• 将跨模态注意力模块集成到类似U-Net的编码器-解码器网络中，用于3D图像体积配准。
• 使用包含卷积层与最大池化层的特征提取器，对输入图像体积进行下采样并提取分层特征。
• 将经过注意力加权的特征送入深度配准器，融合多模态信息并预测空间变换以实现图像对齐。
• 采用基于配准后移动图像与固定图像之间体素级均方误差的损失函数，进行端到端训练。

实验结果

研究问题

• RQ1 自注意力机制是否能通过显式建模空间对应关系，提升MRI与TRUS图像之间的跨模态图像配准性能？
• RQ2 所提出的跨模态注意力模块是否能在远小于标准大型CNN的网络规模下，实现更优的配准性能？
• RQ3 与原始MRI图像相比，使用分割掩码作为输入在配准精度与鲁棒性方面有何影响？
• RQ4 注意力机制在多大程度上通过特征可视化提升了网络的可解释性？
• RQ5 与最先进方法相比，该方法是否能在显著降低推理时间与模型尺寸的前提下，实现更优的性能表现？

主要发现

• 所提出的Attention-Reg网络将目标配准误差（TRE）降低至3.71±1.99 mm，相比基线的10.17±5.75 mm有显著提升。
• 模型仅含1,248,777个参数，优于MSReg（1600万参数）与DVNet（530万参数），在模型尺寸仅为1/10的情况下实现了更优的精度。
• 网络推理时间仅为3ms，为MSReg（6ms）的一半，展现出极高的效率。
• 消融实验表明，移除跨模态注意力模块后性能显著下降（p < 0.001），证实其关键作用。
• 使用分割掩码作为输入可进一步提升精度，使TRE降低至3.60±2.01 mm，并在注意力模块存在时恢复了基于标签输入相对于原始MRI输入的优势。
• Grad-CAM可视化结果表明，注意力机制聚焦于解剖学相关区域（如前列腺边界），证实了有意义的特征学习与空间对齐。