Adaptive Correspondence Scoring for Unsupervised Medical Image Registration

TL;DR

This work tackles the challenge of unsupervised deformable medical image registration when intensity constancy is violated by noise, covisibility issues, and acquisition variability. It introduces an adaptive correspondence scoring framework that learns a dense score map $\hat{S}$ via a scoring network $g_\phi$, and reweights the data-fidelity term during training through $\\mathcal{L}_{de}$, $\\mathcal{L}_{ucs}$, and regularizers $\\mathcal{L}_{reg}$ and $\\mathcal{L}_{smooth}$ with momentum-guided adaptive regularization. The method is trained in a coordinated, alternating fashion with a warm-up phase and is validated across Voxelmorph, Transmorph, and Diffusemorph on three datasets (ACDC, CAMUS, and a private 3D echo), showing statistically significant improvements in Dice, HD, and ASD metrics. The proposed approach achieves robustness to nuisance variability without adding inference-time overhead, highlighting its practical potential for reliable clinical image registration.

Abstract

We propose an adaptive training scheme for unsupervised medical image registration. Existing methods rely on image reconstruction as the primary supervision signal. However, nuisance variables (e.g. noise and covisibility), violation of the Lambertian assumption in physical waves (e.g. ultrasound), and inconsistent image acquisition can all cause a loss of correspondence between medical images. As the unsupervised learning scheme relies on intensity constancy between images to establish correspondence for reconstruction, this introduces spurious error residuals that are not modeled by the typical training objective. To mitigate this, we propose an adaptive framework that re-weights the error residuals with a correspondence scoring map during training, preventing the parametric displacement estimator from drifting away due to noisy gradients, which leads to performance degradation. To illustrate the versatility and effectiveness of our method, we tested our framework on three representative registration architectures across three medical image datasets along with other baselines. Our adaptive framework consistently outperforms other methods both quantitatively and qualitatively. Paired t-tests show that our improvements are statistically significant. Code available at: \url{https://voldemort108x.github.io/AdaCS/}.

Figures (10)

• Figure 1: Existing approaches assume uniform intensity constancy and covisibility across the entire image; during training, this causes irreconcilable penalties, i.e., regions with large error residuals due to the absence of correspondence as highlighted in the red box. Our proposed approach addresses this by re-weighting error residuals with a predictive correspondence scoring map. By doing so, we get a smoother optimization when we reduce the influence of these outliers, leading to an improved performance.
• Figure 2: Diagram of training pipeline of our proposed adaptive scoring framework. Our proposed framework first estimates displacement $\hat{u}$ from source and target image pair $(I_s,I_t)$. We then apply spatial transform to obtain the warped source image $I_s(x+\hat{u})$. Before computing error residuals, we estimate the correspondence scoring map from the target image $I_t$ and then adaptively weight the error residuals for gradient-based optimization. The detailed training strategy is discussed in \ref{['alg:training']}.
• Figure 3: Qualitative evaluation of our method against the second-best approach in each dataset (top two rows: ACDC bernard_deep_2018 and bottom two rows: CAMUS kim_diffusemorph_2022). Each block, delineated by black solid lines, contains source and target images with myocardium segmentation contours. The top row displays the original images, and the bottom row showcases head-to-head comparison (warped source $I_s(x+\hat{u})$) between our method and the second-best method. The yellow highlights indicate the ground truth ES myocardium. Dice scores are reported in the subtitles.
• Figure 4: Qualitative visualization of our proposed framework in Voxelmorph architecture balakrishnan_voxelmorph_2019 on ACDC bernard_deep_2018 (top row) and CAMUS leclerc_deep_2019 (bottom row) validation sets. The third column exhibits successful matching corroborated by the estimated displacement in the fourth column, but the error map in the fifth column reveals residuals. Our predicted scoring map in the sixth column identifies and prevents drift of $f_\theta(\cdot)$, as demonstrated by the re-weighted error in the last column.
• Figure 5: Training dynamics of $\mu_T$ and $m_T$. Left: Mean of error residuals $\mu_T$ in \ref{['eq:error_residuals']}. Right: Adaptive momentum guided weight $m_T$ in \ref{['eq:momentum']}.