[论文解读] Test-Time Training for Deformable Multi-Scale Image Registration

Wentao Zhu, Yufang Huang+4 more|arXiv (Cornell University)|Mar 25, 2021

Medical Image Segmentation Techniques参考文献 46被引用 4

一句话总结

该论文提出了一种基于多尺度深度网络的推理时训练框架，用于可变形图像配准，在噪声大、形变严重的医学图像上提升了泛化能力。通过在推理过程中迭代地在不同尺度上优化配准场，该方法在分割和密集组织追踪任务中优于VoxelMorph和ANTs，且在超声序列上的推理速度比ANTs快20倍。

ABSTRACT

Registration is a fundamental task in medical robotics and is often a crucial step for many downstream tasks such as motion analysis, intra-operative tracking and image segmentation. Popular registration methods such as ANTs and NiftyReg optimize objective functions for each pair of images from scratch, which are time-consuming for 3D and sequential images with complex deformations. Recently, deep learning-based registration approaches such as VoxelMorph have been emerging and achieve competitive performance. In this work, we construct a test-time training for deep deformable image registration to improve the generalization ability of conventional learning-based registration model. We design multi-scale deep networks to consecutively model the residual deformations, which is effective for high variational deformations. Extensive experiments validate the effectiveness of multi-scale deep registration with test-time training based on Dice coefficient for image segmentation and mean square error (MSE), normalized local cross-correlation (NLCC) for tissue dense tracking tasks. Two videos are in https://www.youtube.com/watch?v=NvLrCaqCiAE and https://www.youtube.com/watch?v=pEA6ZmtTNuQ

研究动机与目标

• 解决基于深度学习的图像配准中的领域偏移问题，即模型在噪声大或高度形变的测试图像上泛化能力差。
• 提升无监督深度配准在挑战性临床场景（如超声心动图和3D MR成像）下的鲁棒性和准确性。
• 通过推理时训练实现模型在推理阶段的自适应，缩小训练与测试数据之间的泛化差距。
• 通过从粗到细的多尺度优化策略，处理大而复杂的形变，逐级建模残差形变。
• 在无需真实形变场监督的情况下，为图像分割和密集组织追踪等下游任务实现高精度配准。

提出的方法

• 引入推理时训练，即在每个测试图像对上使用自监督优化对深度网络进行微调，从而减小训练分布与测试分布之间的领域偏移。
• 设计一种多尺度配准框架，在多个分辨率层级上估计残差形变场，从粗到细逐步优化，以稳定训练过程并提升精度。
• 采用U-Net架构在每个尺度上预测残差形变场，前一尺度的输出作为下一尺度的输入，实现逐步精细化。
• 在推理时优化过程中，使用Dice系数、MSE和归一化局部互相关（NLCC）等相似性度量作为自监督信号。
• 通过网络架构和优化过程隐式集成微分同胚约束，确保最终形变场的拓扑一致性。
• 在多样化的医学图像上进行端到端训练，随后仅使用移动图像与固定图像对进行推理时自适应，无需真实标签。

Figure 1: Framework of multi-scale deep registration networks with test-time training in scale $s$ .

实验结果

研究问题

• RQ1 推理时训练是否能显著提升基于深度学习的图像配准在噪声大、高度形变的医学图像上的泛化能力？
• RQ2 多尺度残差形变建模如何提升在大形变情况下的配准精度与稳定性？
• RQ3 所提出的框架是否在分割和密集追踪任务中均优于现有的无监督学习方法（如VoxelMorph）和传统方法（如ANTs）？
• RQ4 推理时训练的计算成本如何？是否可实现临床实用化？
• RQ5 从粗到细的多尺度策略是否能降低复杂配准场景下的优化难度并提升收敛性？

主要发现

• 所提出的推理时训练与多尺度配准方法（SSMSR）显著提升了3D MR图像的配准精度，在分割任务中获得的Dice分数高于VoxelMorph和ANTs。
• 在超声心动图序列中，SSMSR生成的形变场比VoxelMorph和ANTs更平滑且更具解剖学合理性，血流追踪中可清晰识别出涡旋模式。
• 该方法相比ANTs实现20倍提速，单张NVIDIA 1080 Ti GPU上处理49帧超声序列的总推理时间不足九分钟。
• 推理时训练有效缩小了训练与测试性能之间的精度差距，尤其在噪声大、形变严重的图像上表现显著，经定性和定量对比验证。
• 多尺度残差方法能更好地建模大尺度与细尺度的速度变化，体现在心肌运动和血流模式的可视化效果更优。
• 在均方误差（MSE）和归一化局部互相关（NLCC）指标上，该框架在密集组织追踪任务中始终优于VoxelMorph和ANTs。

Figure 2: Reconstruction comparison for cardiac blood from test-time training based multi-scale registration (scale 1/8) w/o (first), i.e., VoxelMorph [ 34 ] , w/ (middle) test-time training and ground truth image (last column).