[论文解读] Attention for Image Registration (AiR): an unsupervised Transformer approach

Zihao Wang, Hervé Delingette|arXiv (Cornell University)|May 5, 2021

Advanced Neural Network Applications参考文献 17被引用 4

一句话总结

本文提出AiR，一种基于Transformer的无监督框架，用于可变形图像配准，将图像配准视为序列到序列的翻译任务。通过利用无需卷积主干网络的多尺度自注意力机制，AiR在基准数据集上实现了最先进性能，在形变准确性和结构相似性方面优于单尺度Transformer和基于CNN的方法，尤其在大形变情况下表现更优。

ABSTRACT

Image registration is a crucial task in signal processing, but it often encounters issues with stability and efficiency. Non-learning registration approaches rely on optimizing similarity metrics between fixed and moving images, which can be expensive in terms of time and space complexity. This problem can be exacerbated when the images are large or there are significant deformations between them. Recently, deep learning, specifically convolutional neural network (CNN)-based methods, have been explored as an effective solution to the weaknesses of non-learning approaches. To further advance learning approaches in image registration, we introduce an attention mechanism in the deformable image registration problem. Our proposed approach is based on a Transformer framework called AiR, which can be efficiently trained on GPGPU devices. We treat the image registration problem as a language translation task and use the Transformer to learn the deformation field. The method learns an unsupervised generated deformation map and is tested on two benchmark datasets. In summary, our approach shows promising effectiveness in addressing stability and efficiency issues in image registration tasks. The source code of AiR is available on Github.

研究动机与目标

• 解决基于相似性度量优化的传统非学习型图像配准方法存在的不稳定性与高计算成本问题。
• 克服基于CNN的深度学习方法在捕捉长距离依赖关系和全局结构形变方面的局限性。
• 探索Transformer架构（此前在自然语言处理中占主导地位）在无需真实形变场监督下的端到端无监督图像配准中的潜力。
• 开发一种多尺度注意力机制，以增强Transformer中跨层级的特征表示，从而更好地建模复杂形变。
• 证明自注意力机制能够以弱监督方式有效学习形变场，避免对标注训练数据的需求。

提出的方法

• 该框架将图像配准视为序列到序列的翻译任务，使用固定图像和移动图像的图像块作为输入标记。
• Transformer编码器通过多头自注意力机制处理固定图像块，学习源图像的记忆表征。
• Transformer解码器结合编码后的特征表示与对移动图像块的自注意力机制，预测形变场。
• 通过可微分空间采样将形变场用于图像形变，使重采样操作支持反向传播。
• 引入多尺度并行注意力（MAPT）模块，以增强不同感受野尺寸下的特征表示，提升在大形变情况下的性能。
• 模型通过基于形变后移动图像与固定图像之间像素级L2距离的重建损失，在无监督设置下进行训练。

实验结果

研究问题

• RQ1 纯Transformer架构是否能在无卷积特征提取器的情况下有效学习图像配准的形变场？
• RQ2 在建模复杂形变时，单尺度Transformer与多尺度注意力机制的性能相比如何？
• RQ3 无监督的Transformer方法是否能在无真实形变标签的情况下实现具有竞争力的配准精度？
• RQ4 Transformer中的自注意力机制是否在捕捉图像配准任务中的全局结构关系方面优于CNN？
• RQ5 多尺度特征学习对生成形变图的质量和图像重建质量有何影响？

主要发现

• 多尺度并行注意力Transformer（MAPT-AiR）在MNIST数据集的所有指标上均表现最佳，平均PSNR为17.311 ± 2.522，Smooth DICE得分为0.827 ± 0.034。
• 单尺度AiR在PSNR方面优于基于CNN的DIRNets（16.966 ± 2.897 vs. 12.941 ± 2.218），但表现不及多尺度变体。
• MAPT-AiR成功捕捉了如数字'9'、'7'和'8'中的大形变，而Single-AiR无法对齐此类复杂形状。
• 模型在MNIST上的平均MSE为0.027 ± 0.015，表明其重建保真度高，像素级误差极小。
• 尽管结构对齐良好，所有基于Transformer的方法在形变图像中均表现出几何细节缺陷和噪声，表明其在像素级精度方面仍存在局限。
• 结果证实，尽管Transformer在全局形变建模方面有效，但在细粒度纹理和边缘保持方面仍落后于CNN。