[论文解读] Affine Medical Image Registration with Coarse-to-Fine Vision Transformer

Tony C. W. Mok, Albert C. S. Chung|arXiv (Cornell University)|Mar 29, 2022

Medical Image Segmentation Techniques被引用 5

一句话总结

本文提出C2FViT，一种用于3D仿射医学图像配准的新型粗到细视觉Transformer方法，通过利用全局自注意力机制和多分辨率特征学习，相较于基于CNN的方法，显著提升了准确性、鲁棒性和泛化能力。其性能与ANTs和Elastix等传统方法相当，同时保持了基于学习方法的快速推理速度。

ABSTRACT

Affine registration is indispensable in a comprehensive medical image registration pipeline. However, only a few studies focus on fast and robust affine registration algorithms. Most of these studies utilize convolutional neural networks (CNNs) to learn joint affine and non-parametric registration, while the standalone performance of the affine subnetwork is less explored. Moreover, existing CNN-based affine registration approaches focus either on the local misalignment or the global orientation and position of the input to predict the affine transformation matrix, which are sensitive to spatial initialization and exhibit limited generalizability apart from the training dataset. In this paper, we present a fast and robust learning-based algorithm, Coarse-to-Fine Vision Transformer (C2FViT), for 3D affine medical image registration. Our method naturally leverages the global connectivity and locality of the convolutional vision transformer and the multi-resolution strategy to learn the global affine registration. We evaluate our method on 3D brain atlas registration and template-matching normalization. Comprehensive results demonstrate that our method is superior to the existing CNNs-based affine registration methods in terms of registration accuracy, robustness and generalizability while preserving the runtime advantage of the learning-based methods. The source code is available at https://github.com/cwmok/C2FViT.

研究动机与目标

• 解决基于CNN的仿射配准方法在处理大初始错位时的局限性，并提升其在未见数据集上的泛化能力。
• 开发一种基于学习的仿射配准框架，以更有效地捕捉全局几何依赖关系，超越局部卷积归纳偏置的影响。
• 提升3D脑图像配准任务中的配准准确性和鲁棒性，特别是在具有挑战性的空间初始化条件下。
• 在保持深度学习方法运行时效率的同时，实现与传统优化方法相当或更优的性能。
• 证明所提出范式在少量架构修改下，对多种参数化配准任务的适应能力。

提出的方法

• C2FViT采用粗到细的架构，通过在多个分辨率下处理图像对，逐步优化仿射变换估计。
• 模型采用卷积视觉Transformer（CvT）主干网络，结合局部归纳偏置与全局自注意力机制，以更好地建模长距离空间依赖关系。
• 多分辨率策略使网络能够在不同尺度上捕捉全局方向和细粒度的空间对齐特征。
• 在推理过程中应用渐进式空间变换，以稳定训练并提升收敛性，尤其在大初始错位条件下表现更优。
• 采用质心（CoM）初始化作为几何先验，以增强鲁棒性并减少对随机空间初始化的依赖。
• 通过解剖学标签引入半监督学习，以增强泛化能力，且不增加模型参数量或计算成本。

实验结果

研究问题

• RQ1 在大初始错位条件下，基于视觉Transformer的架构是否能优于基于CNN的方法在3D仿射配准中的表现？
• RQ2 所提出的C2FViT模型在训练分布之外的未见数据集上泛化能力如何？
• RQ3 全局自注意力机制与多分辨率处理的结合在多大程度上提升了配准的准确性和鲁棒性？
• RQ4 所提出方法是否能在保持与传统优化方法相当性能的同时，维持快速的推理速度？
• RQ5 如渐进式空间变换和CoM初始化等架构组件对模型性能与稳定性有何影响？

主要发现

• 在LPBA数据集上，C2FViT的Dice分数为0.66，HD95为2.96，与ANTs和Elastix等传统方法性能相当。
• 在OASIS数据集上，采用半监督学习的模型Dice分数达到0.69，HD95为2.81，优于无监督和有监督的基线方法。
• C2FViT保持了快速推理，平均每次配准耗时0.09秒，显著快于ANTs和Elastix。
• 消融实验证实，渐进式空间变换和CoM初始化使Dice分数提升+0.02，HD95降低-0.37。
• 与直接矩阵估计相比，解耦的仿射变换模型在精度上略有提升（Dice +0.02），且运行时开销极低。
• 该模型在未见数据集上表现出良好的泛化能力，即使在训练分布外的数据上测试也展现出强鲁棒性。