Adaptive Correspondence Scoring for Unsupervised Medical Image Registration

Deformable medical image registration aims to accurately determine non-rigid correspondences through dense displacement vectors between source and target images. This process is a crucial step for medical image analysis, such as tracking disease progression for diagnosis and treatment [29, 14, 39]. Due to the impracticality of obtaining ground truth displacement, it has been a long-standing problem and has been extensively studied in the past decades [23, 31, 2, 5, 8, 22, 41, 38, 40].

Classical methods approach this challenge by solving an iterative pair-wise optimization problem between source and target images using elastic-type models [11, 23], free-form deformations with b-splines [31], and topology-preserving diffeomorphic models [2, 3]. However, these approaches are computationally expensive and time-consuming, limiting their practical utility in large-scale real-world data. Recently, learning-based approaches have been widely adopted for their speed (GPU accelerated forward pass that is hundreds of times faster in inference speed than classical methods) and state-of-the-art performance [5, 41, 22, 8, 9, 33]. Due to lack of ground truth displacement, these approaches leverage dense image or volumetric reconstruction in an unsupervised setting, or make use of segmentation masks in a weakly supervised setting. To train these approaches via gradient-based optimization, a source image is warped by the estimated displacement to reconstruct a target image. The supervision signal comes from minimizing the reconstruction error between warped source and target as the data-fidelity term, along with a regularizer based on the assumption that the object imaged is locally smooth and connected. The feasibility in minimizing this objective relies on intensity constancy between the two images during imaging.

However, assuming such intensity constancy uniformly across the entire image domain neglects the facts of motion ambiguity and non-uniform noise corruption in medical images. Assuming sufficiently exciting local regions, only a subset of pixels of one image can be uniquely matched or corresponded to another based on the image intensity profiles. As shown in Fig. 1, despite a displacement estimator model predicting largely correct corresponding pixels between the two images, the error residuals between the warped source and target images are still non-zero. In fact, they are dominated by regions with no correspondence; when this phenomenon occurs during the training of a displacement estimator, it results in performance degradation (driven towards an undesirable minimum) in the subsequent optimization steps due to the noisy gradients. To address this issue, we propose an unsupervised correspondence scoring framework that identifies regions with high chance of establishing correspondence and adaptively reduces the influence of error residuals from nuisance variability during training. This prevents the displacement estimator from drifting away due to large residuals caused by the lack of correspondence between input image pairs.

Our proposed scoring estimator is deployed during the unsupervised training of a displacement estimator and yields a soft scoring map. It is optimized alternatingly together with the displacement estimator to adaptively re-weight the data-fidelity term in order to mitigate the negative impact of nuisances when establishing correspondence. We introduce an unsupervised training objective for the scoring estimator to learn an accurate correspondence scoring map without additional annotation. Yet, there exists a trivial solution of predicting a score map of all zeros. Therefore, we additionally propose a regularizer for the scoring estimator to bias it away from the trivial solution, along with a momentum-guided adaptive total variation to encourage smoothness in the scoring map. To validate the effectiveness of our proposed method, we tested on three different datasets including: (1) ACDC [6], a public 2D MRI dataset, (2) CAMUS [24], a public 2D ultrasound dataset, and (3) a private 3D echocardiography dataset [1, 35]. To further show the versatility of our proposed method, we tested on three representative unsupervised image registration architectures including: (1) Voxelmorph [5], (2) Transmorph [8] and (3) Diffusemorph [22] along with other baselines. Our proposed framework can be applied in a plug-and-play manner to consistently improve existing methods. Paired t-tests show that the improvements obtained by utilizing our approach are statistically significant.

Our contributions are as follows: (1) We propose an adaptive framework that incorporates correspondence scoring for unsupervised deformable medical image registration. (2) We introduce an unsupervised correspondence scoring network to be used during the training of a displacement estimator. The scoring network learns to determine whether a given image allows for establishing correspondence by minimizing the typical image reconstruction loss with scoring and momentum-guided adaptive regularization. (3) Our proposed method consistently outperforms other baselines across three representative registration architectures over three medical image datasets with diverse modalities. The performance gain comes with no cost in memory or run time during inference, but only the deployment of our scoring estimation during training.

Unsupervised medical image registration. Balakrishnan et al. [5] proposed an unsupervised learning framework using U-Net as the displacement estimator. This framework imposes intensity constancy by minimizing the mean squared error between the warped source image and the target image to update the parametric displacement estimator via gradient-based optimization. A number of works have been developed upon this architecture including adding diffeomorphic regularization (Voxelmorph-diff) [10], jointly learning amortized hyperparameters (Hypermorph) [18] and learning contrast-invariant registration without acquired images (SynthMorph) [15]. Recently, Zhang et al. [40] proposed an uncertainty estimation framework that extends the widely used homoscedastic assumption in objectives to heteroscedastic assumption. Inspired by the recent advances in vision transformers [25, 12], Chen et al. [8] introduced TransUNet [7], a hybrid Transformer-CNN architecture. This design replaces encoders with Swin Transformer [25] to enhance the receptive field while preserving convolutional decoders to bolster the model’s ability to capture long-range motion. A number of extensions to this architecture have been proposed, such as substituting convolutional decoders with transformer layers [33] and incorporating multi-scale pyramids [26]. Additionally, score-based diffusion models such as DDPM [34] have shown high-quality performance in generative modeling. To leverage the advantage of DDPM, Kim et al. [22] presents a diffusion-based architecture, composed of a diffusion network and a deformation network. Recent work built upon this architecture explores adding feature and score-wise diffusion [30].

We demonstrate our proposed adaptive framework and compare it with with other related formulations as baselines on three representative architectures including: (1) Voxelmorph [5], (2) Transmorph [8] and (3) Diffusemorph [22]. We also tested our proposed framework on c-LapIRN [27] and deferred the results to the Supp. Mat. due to the page limit.

Adaptive weighting schemes. A wide range of image processing problems involve optimizing an energy function that combines a data-fidelity term and a regularization term. The relative importance between the two terms is usually weighted by a scalar, which disregards the heteroscedastic nature of error residuals [36]. To address this challenge, several adaptive weighting schemes are proposed in the spatial domain and over the course of optimization based on the local residual [17, 16, 37]. Wong et al. [36] later provides a data-driven algorithm that deals with multiple frames [36]. Zhang et al. [40] proposed a … In this paper, we selected AdaReg [37] and AdaFrame [36] as our baselines for adaptive weighting. The aforementioned methods are not learning-based; whereas, we proposed a learning-based correspondence scoring in a collaborative framework.

Aleatoric uncertainty estimation. Our proposed adaptive correspondence scoring is conceptually related to aleatoric uncertainty modeling in the Bayesian learning framework, which aims to estimate input-dependent noise inherent in the observations [21, 4, 28, 19, 32]. This can be attributed to for example motion noise or sensor noise, resulting in uncertainty which cannot be reduced even if more data were to be collected. Kendall et al. [21] proposed a maximum likelihood estimation (MLE) formulation that minimizes the negative log-likelihood (NLL) criterion using stochastic gradient descent. This approach re-weights the data-fidelity term using predictive variance estimates mediated by a regularization term after assuming noise distribution is heteroscedastic Gaussian. Seitzer et al. [32] later identifies that such formulation using inverse variance weighting will result in over-confident of variance estimates, leading to undesired undersampling. Thus, an exponentiated $\beta$ term is proposed in the new loss formulation, termed $\beta$-NLL, to counteract the undersampling leading to undesired performance. In this paper, we selected NLL [21] and $\beta$-NLL [32] as baselines for aleatoric uncertainty estimation with a U-Net based variance estimator that is jointly trained with displacement estimator under each formulation.

Let $I_{s}$ be the source image and $I_{t}$ be the target image, where $I:\Omega\mapsto[0,1]$ is the imaging function after normalization and $\Omega$ is the image space. Deformable image registration aims to estimate a dense displacement vector field that characterizes the correspondence between two images for each pixel $\hat{u}:\Omega\mapsto\mathbbm{R}^{2}$ ($\Omega\mapsto\mathbbm{R}^{3}$ for 3D data). For ease of notation and terminology, we will use the terms images and pixels to refer to both 2D and 3D data.

Due to the lack of ground truth, intensity constancy and the smoothness assumption are imposed to constrain the parametric model $f_{\theta}(\cdot)$ (e.g., a neural network) as displacement estimator $\hat{u}=f_{\theta}(I_{t},I_{s})$ by minimizing the following objective to update the parameters $\theta$ in the model

where $x\in\Omega$ denotes a coordinate and $\lambda$ denotes the hyperparameter to modulate the trade-off between two terms. In this work, we explore three representative deep neural network architectures that serve as displacement estimators including (1) convolution-based (Voxelmorph [5]), (2) transformer-based (Transmorph [8]) and (3) diffusion-based (Diffusemorph [22]).

Datasets. We tested our proposed framework on three different cardiac datasets, including two 2D public datasets, containing two medical imaging modalities (MRI and ultrasound), and one private 3D dataset. To construct the image pair, we selected the end-diastole (ED) frame as the source image and the end-systole (ES) frame as the target image. ED to ES registration is considered long-range and most challenging in the cardiac sequence. The detailed steps of dataset preprocessing can be found in the Supp. Mat.

ACDC [6]. The ACDC dataset contains 2D human cardiac MRI from 150 patients with various cardiac conditions. We randomly selected 80 patients containing 751 image pairs for training, 20 patients containing 200 pairs for validation, and the remaining 50 patients containing 538 pairs for testing.

CAMUS [24]. The CAMUS dataset contains 2D human cardiac ultrasound images from 500 subjects. We randomly selected 600 image pairs for training, 200 pairs for validation, and another 200 pairs for testing.

Private 3D Echocardiography [35, 1]. To validate the effectiveness of our proposed method, we also tested on a private 3D echocardiography dataset and reported our results in the Supp. Mat.

Evaluation metrics. We evaluate our results quantitatively by warping myocardium segmentation in the source image with our predicted displacement $\hat{u}$ and compute anatomical conformance in terms of (1) Dice coefficient score (DSC) (2) Hausdorff distance (HD) and (3) Average surface distance (ASD) with ground truth segmentation in the target image. Definitions of metrics can be found in the Supp. Mat.

Implementation details. All our experiments were implemented using Pytorch on NVIDIA V100/A5000 GPUs. The architecture of the scoring estimator is implemented on a U-Net backbone. Code is provided to ensure reproducibility. To show the versatility of our proposed framework, we tested on three representative unsupervised registration architectures for each dataset: (1) Voxelmorph [5], (2) Transmorph [8] and (3) Diffusemorph [22]. We also tested on c-LapIRN architecture [27] and present the quantitative results in the Supp. Mat. due to the page limit. The descriptions of baselines are summarized below and remaining details (i.e., hyperparameters) can be found in the Supp. Mat.

Adaptive weighting schemes. We present two baselines AdaReg and AdaFrame that, like us, utilize adaptive weighting schemes during training.

AdaReg [37]. We compute first the local error residual $\rho=|I_{t}(x)-I_{s}(x+\hat{u})|$ and the global residual $\sigma=\frac{1}{\frac{1}{|\Omega|}\sum_{x\in\Omega}|I_{t}(x)-I_{s}(x+\hat{u})|}$ using the displacement estimator prediction $\hat{u}$. We then compute the adaptive regularization weighting as $\alpha(x)=\exp(-\frac{c\rho}{\sigma})$, where $c=50$. We then optimize the displacement estimator using the loss $\mathcal{L}_{\text{AdaReg}}=\frac{1}{|\Omega|}\sum_{x\in\Omega}[I_{t}(x)-I_{s}(x+\hat{u})]^{2}+\lambda\|\alpha(x)\nabla\hat{u}\|^{2}$.

AdaFrame [36]. We first compute the local error residual $\delta=|I_{t}(x)-I_{s}(x+\hat{u})|$ and then normalize it with its mean $\mu$ and standard deviation $\sigma$ as $\rho=\frac{\delta-\mu}{\sqrt{\sigma^{2}+\epsilon}}$. We then compute the adaptive weight activated by a scaled and shifted sigmoid function as $\alpha(x)=1-\frac{1}{1+\exp(-(a\rho-b))}$ where $a=\frac{a_{0}}{\mu+\epsilon}$ and $b=b_{0}(1-\cos\pi\mu)$. We choose $a_{0}=0.1$ and $b_{0}=10$. We then optimize the displacement estimator using the loss $\mathcal{L}_{\text{AdaFrame}}=\frac{1}{|\Omega|}\sum_{x\in\Omega}\alpha(x)[I_{t}(x)-I_{s}(x+\hat{u})]^{2}+\lambda\|\nabla\hat{u}\|^{2}$.

Aleatoric uncertainty estimation. We present two baselines that utilize aleatoric uncertainty estimates in terms of predictive variance to weight the objective during training adaptively. In order to obtain the predictive variance, we utilize a U-Net as variance estimator $h(\cdot)$ that takes target image $I_{t}$ and warped source $I_{s}(x+\hat{u})$ as input to predict the noise variance $\hat{\sigma}_{I}=h(I_{t},I_{s}(x+\hat{u}))$. The variance estimator is trained jointly with the displacement estimator.

NLL [21]. To jointly train the displacement and variance estimators, we compute the loss as $\mathcal{L}_{\text{NLL}}=\frac{1}{|\Omega|}\sum_{x\in\Omega}\frac{1}{\hat{\sigma}_{I}^{2}(x)}[I_{t}(x)-I_{s}(x+\hat{u})]^{2}+\log\hat{\sigma}^{2}_{I}(x)$

$\beta$-NLL [32]. To jointly train the displacement and variance estimators under $\beta$-NLL objective, we compute the loss as $\mathcal{L}_{\beta\text{-NLL}}=\frac{1}{|\Omega|}\sum_{x\in\Omega}\lfloor\hat{\sigma}_{I}(x)^{2\beta}\rfloor(\frac{1}{\hat{\sigma}_{I}^{2}(x)}[I_{t}(x)-I_{s}(x+\hat{u})]^{2}+\log\hat{\sigma}^{2}_{I}(x))$ where $\beta=0.5$.

Registration accuracy. We present our quantitative evaluation in Tab. 1, where our proposed method shows consistent improvement over other baselines. Surprisingly, we observe that incorporating baselines (adaptive weighting schemes and uncertainty estimation) tend to harm performance. Unlike them, our method improves over the base models. Compared to the base models (Voxelmorph [5], Transmorph [8], and Diffusemorph [22]), which are the second-best methods in each architecture, our proposed method performs better especially in terms of Dice score on each dataset. We additionally conducted a paired t-test of our proposed method to show that our consistent improvement is statistically significant as in Tab. 3.

To quantitatively evaluate the registration performance, we plot the warped source image along with segmentation overlayed with ground truth shown in Fig. 3. Our proposed framework is visually better across datasets with different modalities and frameworks. Both the quantitative and qualitative evaluations show that our proposed framework captures more accurate correspondence, validating the effectiveness of our proposed adaptive scoring.

Tab. 1 shows that both uncertainty weighting strategies, NLL [21] and $\beta$-NLL [32], did not yield improvements over the baseline. This observation indicates that optimization of uncertainty estimation cannot be approached with a simple joint optimization alongside image registration. Furthermore, we observe that adaptive regularization techniques, exemplified by AdaFrame [36] and AdaReg [37], were ineffective in enhancing performance, which could be attributed to the computation of adaptive weights based on statistical assumptions of error residuals contrary to our proposed unsupervised learning approach. We note a discernible decline in performance when transitioning from Voxelmorph [5], which utilizes ConvNets, to Transmorph [8] and Diffusemorph [22]. This decline underscores the challenges faced by transformer and diffusion-based models when applied to datasets comprising smaller-scale medical images. We also present some failure cases in the Supp. Mat., where the myocardium (typically challenging due to irregular volumes) is considerably thin.

Results of correspondence scoring map and adaptive weighting. To qualitatively evaluate the effectiveness of our proposed correspondence scoring during training, we present Fig. 4 to show that our predicted scoring map accurately identifies the regions with low correspondence and prevents the displacement estimator $f_{\theta}(\cdot)$ from drifting away by re-weighting the error residuals.

In Fig. 5, we illustrate the behavior of our proposed adaptive momentum-guided weight, denoted as $m_{T}$ (Eq. 7). This weight dynamically adapts during training, responding to the evolving characteristics of error residuals (Eq. 5). The observed evolution of our adaptive weight aligns with our design objective, progressively enhancing smoothness penalization as training toward convergence. The effectiveness of $m_{T}$ is also evidenced by our quantitative ablation analysis in Fig. 5, in which our proposed formulation generally achieves a better performance with the adaptive weight.

Ablation study. We present Tab. 2 for ablation studies to show the effectiveness of each loss term. Without $\mathcal{L}_{\text{reg}}$ and $\mathcal{L}_{\text{smooth}}$, the predicted correspondence scoring map $\hat{S}$ will result in a degenerate solution of all zeros. By introducing regularization and smoothness constraints, the performance of our proposed framework steadily increases across various datasets and registration architectures. We also perform an ablation study comparing with robust losses including mutual information (MI) and normalized cross correlation (NCC) shown in Tab. 4 in the Supp. Mat. Our conclusions still hold as our method yields consistent performance improvements over the baselines.

Training stability and sensitivity to hyperparameters. We first show the training stability by plotting the training curves of our proposed unsupervised scoring term $\mathcal{L}_{\text{ucs}}$ in Eq. 3 and regularization term $\mathcal{L}_{\text{reg}}$ in Eq. 4 in the left figure of Fig. 6, demonstrating both terms can be jointly minimized. We further conduct a sensitivity study on hyperparameters $\alpha$ and $\beta$ in Eq. 9. While they are not very sensitive (change within $\approx$ 1% Dice for various $\alpha,\beta$), tuning more can indeed yield better results.

Evaluation on smoothness and training time. We provide standard deviation and measure of smoothness (percentage of estimated displacement $\hat{u}$ with a negative Jacobian determinant, i.e., $|J_{\hat{u}}|\leq 0$) in Table 5. All methods have $<$1%, implying physical plausibility and smoothness. Our method produces smoother deformations in Voxelmorph and Transmorph and achieves similarly smooth deformations in Diffusemorph, with $<$0.1%.

In this paper, we propose an adaptive correspondence scoring framework for unsupervised image registration to prevent the displacement estimator from drifting away by noisy gradients caused by low correspondence due to the nuisance variables such as noise or covisibility during training. We introduce a unsupervised correspondence scoring estimation scheme with both scoring and momentum-guided adaptive regularizations to prevent the scoring estimator from a degenerate solution and ensure scoring map smoothness. We demonstrate the effectiveness of our proposed framework on three representative registration architectures and we show consistent improvement compared with other baselines across three medical image datasets with diverse modalities. Though our proposed framework is promising, the hyperparameters in the scoring estimator loss does need to be tuned for each architecture on each dataset. Nonetheless, performance is not too sensitive to hyperparameter tuning, so one does not need to tune meticulously. In the future, we aim to explore an amortized hyperparameter optimization scheme during training to reduce the computation and validate our proposed framework on clinical datasets for further impact.

This work is supported by NIH/NHLBI grant R01HL121226.