[Paper Review] Test-Time Training for Deformable Multi-Scale Image Registration

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

Medical Image Segmentation Techniques46 references4 citations

TL;DR

This paper proposes a test-time training framework with multi-scale deep networks for deformable image registration, improving generalization on noisy, highly deformed medical images. By iteratively refining registration fields across scales during inference, it achieves superior accuracy over VoxelMorph and ANTs in segmentation and dense tracking tasks, with a 20x speedup over ANTs on ultrasound sequences.

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

Motivation & Objective

• Address the domain shift problem in deep learning-based image registration, where models generalize poorly on noisy or highly deformed test images.
• Improve the robustness and accuracy of unsupervised deep registration for challenging clinical scenarios like echocardiography and 3D MR imaging.
• Reduce the generalization gap between training and test data by enabling model adaptation during inference via test-time training.
• Handle large and complex deformations through a coarse-to-fine multi-scale optimization strategy that models residual deformations sequentially.
• Achieve high-precision registration for downstream tasks such as image segmentation and dense tissue tracking without requiring ground-truth deformation fields.

Proposed method

• Introduce test-time training where the deep network is fine-tuned on each test image pair using self-supervised optimization, reducing the domain shift between training and test distributions.
• Design a multi-scale registration framework that estimates residual deformation fields at multiple resolution levels, starting from coarse to fine, to stabilize optimization and improve accuracy.
• Use a U-Net architecture to predict residual deformation fields at each scale, with the output of one scale serving as input to the next for progressive refinement.
• Apply similarity metrics such as Dice coefficient, MSE, and normalized local cross-correlation (NLCC) as self-supervision signals during test-time optimization.
• Integrate diffeomorphic constraints implicitly through the network architecture and optimization process to ensure topological consistency in the final deformation fields.
• Perform end-to-end training on a diverse set of medical images, followed by test-time adaptation using only the moving and fixed image pair, without requiring ground-truth labels.

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

Experimental results

Research questions

• RQ1 Can test-time training significantly improve the generalization of deep learning-based image registration on noisy and highly deformed medical images?
• RQ2 How does multi-scale residual deformation modeling enhance the accuracy and stability of registration in the presence of large deformations?
• RQ3 Does the proposed framework outperform existing unsupervised learning-based methods like VoxelMorph and traditional methods like ANTs in both segmentation and dense tracking tasks?
• RQ4 What is the computational cost of test-time training, and can it be made practical for clinical use?
• RQ5 Can the coarse-to-fine multi-scale strategy reduce optimization difficulties and improve convergence in complex registration scenarios?

Key findings

• The proposed test-time training with multi-scale registration (SSMSR) significantly improves registration accuracy on 3D MR images, achieving higher Dice scores in segmentation tasks compared to VoxelMorph and ANTs.
• On echocardiographic sequences, SSMSR produces smoother and more anatomically plausible deformation fields than VoxelMorph and ANTs, with clearly identifiable vortex patterns in blood flow tracking.
• The method achieves a 20x speedup over ANTs, with a total test-time of under nine minutes for a 49-frame ultrasound sequence on a single NVIDIA 1080 Ti GPU.
• Test-time training reduces the accuracy gap between training and test performance, particularly on noisy images with large deformations, as shown in qualitative and quantitative comparisons.
• The multi-scale residual approach enables better modeling of both large-scale and fine-scale velocity variations, as evidenced by improved visualization of myocardial motion and blood flow patterns.
• The framework consistently outperforms VoxelMorph and ANTs in both mean square error (MSE) and normalized local cross-correlation (NLCC) metrics for dense tissue tracking.

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).