Source-linked AI summary
Recursive Uncertainty-Gated Image Registration for Learning-based Algorithms
Clara Rodrigo González, Oscar Bates, Fu Siong Ng, Meng-Xing Tang
TL;DR
Image correspondence is spatially uneven, and learned registration can struggle with challenging or out-of-domain image pairs. RUGI iteratively refines deformation fields with spatial gating, improving registration across cardiac MRI and ultrasound and reducing pretrained-model MSE by 27–37%.
Problem
Image correspondence is spatially ill-posed and uneven across regions, while learned registration can struggle with out-of-domain or challenging samples.
Source: §1 · pp. 1–2Method
RUGI iteratively predicts deformation updates and gates them spatially, using uncertainty or image residual error to focus refinement.
Source: §1 · p. 2, §3.1 · p. 3Results
RUGI significantly improves registration over single-step inference across cardiac MRI and ultrasound, while error-gated pretrained VoxelMorph, TransMorph, and CycleMorph achieve MSE reductions of 27–37%.
Source: §5 · p. 16, §6 · p. 17Takeaways & Limitations
RUGI’s uncertainty maps identify spatial variation in registration reliability, and improved registration is reflected in downstream functional-parameter estimation.
Source: §6 · p. 17Iterative inference increases computational cost relative to a single pass and may be prohibitive for real-time applications.
Source: §5 · p. 16
Abstract
from arXiv · showhide
Conventional image registration algorithms are robust to domain shifts and achieve low errors, but they are slow and computationally expensive. Deep-learning methods are efficient at inference-time, but face challenges in out-of-domain samples. We propose Recursive Uncertainty-Gated Image Registration (RUGI), an algorithm for iteratively refining deformation fields predicted by learning-based registration models. At each iteration, the registration model predicts an incremental deformation, and a gating map modulates the update. Refinements are hence concentrated in regions that remain difficult to register. We explore two gating strategies: a learned uncertainty-based approach and an image residual error approach. We evaluate RUGI on cardiac MRI and echocardiography datasets and show consistent improvements over single-step inference. Ablation experiments demonstrate that iterative refinement alone improves registration, but informative spatial gating provides a significant additional benefit. The error-gated variant of RUGI can also be applied directly to existing pretrained models; applied to VoxelMorph, TransMorph, and CycleMorph, it yields MSE reductions of 27-37% with no modification to the original training procedure. The improvements in registration performance are reflected in decreased errors in ejection fraction estimation relative to ground truths. These results demonstrate that spatially selective iterative refinement provides an effective strategy to improve registration accuracy at inference-time.
1 Introduction
Image registration supports medical imaging tasks, but learned methods’ speed does not remove their difficulty with out-of-domain or spatially challenging pairs. RUGI addresses this with gated iterative refinement, evaluates uncertainty- and error-based gates, and validates uncertainty and cardiac functional parameter estimation.
1 Introduction
- Spatially varying ill-posedness means image regions differ in how well they constrain correspondence, limiting fixed models’ reliable generalization to new domains.Clear structure and simple motion tend to yield stable deformations, while low-quality images or complex deformations leave more possible correspondences.Source: §1 · pp. 1–2
- Conventional registration is robust to domain shifts but computationally expensive, whereas learning-based methods are fast yet struggle with out-of-domain or challenging samples.Medical image registration supports diagnostics, treatment planning, intervention, longitudinal analysis, multimodal fusion, and ultrasound motion compensation.Source: §1 · pp. 1–2
- RUGI introduces gated iterative updates that concentrate refinement in regions needing it.The framework evaluates predictive-uncertainty and residual-image-error gates, with both improving accuracy over single-step inference.Source: §1 · p. 2
- An ablation study separates the contributions of iterative refinement and spatial gating.Source: §1 · p. 2
- RUGI can be applied to pretrained models without retraining, and the authors validate uncertainty estimation and cardiac functional parameter estimation.Source: §1 · p. 2
2 Background
Prior iterative and test-time registration methods refine predictions but generally apply updates across the image and may require architectural changes or retraining. Prior uncertainty work estimates or uses uncertainty in other ways, leaving inference-time uncertainty-guided selective refinement as the stated gap.
2 Background
- Cascaded, coarse-to-fine, and recurrent methods improve alignment through sequential predictions, but refine uniformly and may require retraining after architectural changes.Examples include LapIRN [3], cascaded VoxelMorph approaches, and recursive or recurrent networks.Source: §2 · p. 2
- Test-time optimization can improve accuracy, but prior approaches modify model architecture and apply refinement without spatial selectivity.The cited approaches include residual-deformation fine-tuning [8], meta-learning for test-time optimization [9], and iterative multi-scale inference.Source: §2 · p. 2
- Prior uncertainty methods weight regularization or consistency constraints [11], while an inference-time framework estimates uncertainty for pretrained models.The latter relates estimated uncertainty to registration error across anatomical structures and pretrained models, but does not apply it to improve registration itself.Source: §2 · p. 2
- The stated gap is inference-time use of uncertainty to selectively refine deformation fields without retraining.Source: §2 · pp. 2–3
3 Methods
RUGI starts from a learned deformation prediction and repeatedly refines the warped image, using a spatial gate to modulate each incremental update. The gate can come from a learned uncertainty estimate or from image residual error.
3 Methods
- Unsupervised registration learns a deformation field by minimizing a loss between the warped moving image and fixed image, then predicts a single field at inference.Source: §3 · p. 3
3.1 Recursive Uncertainty-Gated Image Registration (RUGI)
- At each RUGI iteration, the registration model predicts a dense deformation from the fixed and latest warped images, while a normalized spatial gate modulates the update.The updates are smoothed and normalized; iterations stop according to image similarity, with a Jacobian-based safeguard against severe folding rather than strict diffeomorphism.Source: §3.1 · p. 3, Figure 1 · p. 4
3.1.1 Uncertainty-based gating
- Uncertainty gating gives smaller updates to low-uncertainty regions and continues refining high-uncertainty regions, using a U-Net trained against per-pixel deformation-field standard deviation.The uncertainty map is smoothed and normalized before gating; the added network is described as small, with 465,138 parameters.Source: §3.1.1 · p. 4
3.1.2 Error-based gating
- Error gating smooths and normalizes per-pixel residual error, assigning larger updates to greater mismatches and requiring no retraining when datasets change.It uses only the registration model and allows the residual measure to be replaced; across-iteration deformation spread is a heterogeneity measure, not calibrated uncertainty.Source: §3.1.2 · pp. 4–5
3.2 Training Procedure
The uncertainty network is trained alongside the registration model using sampled deformation variability as its supervision. At inference, its predictions are smoothed and normalized into the gating map.
3.2 Training Procedure
- The registration network is sampled five times per training batch, and the per-pixel standard deviation of its deformation fields supervises the uncertainty network with MSE.The registration network is trained using image-similarity MSE plus bending-energy regularization; the uncertainty network learns predictions specific to the registration model.Source: §3.2 · p. 5, Figure 2 · p. 5
- The uncertainty network estimates epistemic uncertainty: high values mark regions where predictions vary across passes, while consistently registered background or static anatomy has low uncertainty.Source: §3.2 · p. 5
3.3 Implementation details
The implementation uses a compact registration network, a lightweight uncertainty model, and local-NCC residual error to guide selective refinement. It evaluates the method on cardiac MRI and ultrasound datasets.
3.3 Implementation details
- The registration model is a compact Bayesian U-Net-inspired network with two downsampling stages, eight convolutional layers, and approximately 1.9M trainable parameters.Images are resized to 128x128 pixels; at inference, stochasticity is turned off.Source: §3.3 · p. 6
- A 465,138-parameter, two-layer U-Net approximates empirical uncertainty and is designed to be recomputed at each refinement step using only the image pair.The authors selected it to keep uncertainty estimation quick, with cost comparable to or below registration itself.Source: §3.3 · p. 6
- The error-based gating map uses local normalised cross-correlation, followed by an exponential transform to amplify contrast between high- and low-error regions.The similarity measure was chosen for robustness across modalities and insensitivity to global intensity differences.Source: §3.3 · p. 6
- Evaluation uses ACDC cardiac MRI and CAMUS cardiac ultrasound data, with patient-level splits that keep each patient in only one split.ACDC training comprised 100 patients and 560 image pairs; CAMUS included 450 patients and 4,500 training images.Source: §3.3 · p. 6
3.4 Ablation Study
The ablation study compares the full method with variants that remove gating, replace it with random gating, or use a smoothed intensity-based map. These comparisons assess the roles of iterative registration and spatial guidance.
3.4 Ablation Study
- The full method combines uncertainty-based and error-based gating and is compared with three ablations.The ablations test iterative refinement without gating, spatially varying gating, and whether a simpler structural prior can guide refinement.Source: §3.4 · p. 6
- A1 removes spatial gating and applies the full predicted update at each iteration, testing improvement from iterative registration alone.The gating map neither modulates nor spatially selects the update in this variant.Source: §3.4 · p. 6
- A2 replaces the gating map with a smoothed and normalised random map to test the effect of spatially varying gating.Source: §3.4 · p. 6
- A3 uses a Gaussian-blurred, normalised moving image as its gating map to test whether a simpler image-based structural prior can guide refinement.Source: §3.4 · p. 6
3.5 Application to pretrained models
The authors apply error-based iterative gating to pretrained registration models without changing their original training procedure. The experiment compares each model's one-step output with its output in the refinement framework.
3.5 Application to pretrained models
- Error-based gating can be used as an inference-time plug-in iterative refinement framework without modifying the original training procedure.Source: §3.5 · p. 6
- The method was applied to pretrained TransMorph, VoxelMorph, and CycleMorph models whose weights were extracted from [19].Source: §3.5 · p. 6
- The models were trained on the IXI 3-D brain MRI registration dataset, and their original one-step outputs were compared with iterative error-gated outputs.Source: §3.5 · p. 7
- The experiment was designed to assess accuracy gains without retraining and generalisation across architectures, using the same metrics as the main experiments.Source: §3.5 · p. 7
3.6 Evaluation
The evaluation uses patient-level image-registration metrics and statistical tests on held-out cohorts. A separate analysis measures spatial overlap between predicted uncertainty and residual-error regions.
3.6 Evaluation
- Held-out test cohorts comprised 50 ACDC patients and 100 CAMUS patients, with splits by patient to prevent overlap between training and test cohorts.Source: §3.6 · p. 7
- Registration accuracy was assessed using MSE, NCC, MI, and SSIM, with patient-level aggregation, two-sided Wilcoxon signed-rank tests, and BH correction.Effect sizes used paired rank-biserial correlation, with positive differences favouring RUGI; wpBH < 0.05 was considered significant.Source: §3.6 · p. 7
- The ablation study combines similarity metrics using mean signed-rank effect G and assesses significance with a patient-level sign-flip permutation test using 100,000 permutations and BH correction.Source: §3.6 · p. 7
3.7 Validation of Estimated Uncertainty
- Uncertainty validation compares the top 20% of predicted-uncertainty pixels with the top 20% of residual-error pixels; random correspondence has an expected overlap of 20%.Tolerance-aware overlap additionally dilates the masks by 1, 2, 3, or 5 pixels to account for small spatial offsets.Source: §3.7 · p. 7
3.8 Functional Parameters
The section estimates ejection fraction from registered end-systolic and end-diastolic left-ventricular masks, then uses uncertainty-scaled deformation perturbations to quantify uncertainty in functional parameters. With 2D masks, the EF measure is an area-based surrogate rather than clinical volumetric EF.
3.8 Functional Parameters
- Registration estimates motion from end-systole to end-diastole, warps the end-systolic left-ventricular mask, and computes EF from the relative chamber-size change.With perfect registration, the warped end-systolic area approaches the end-diastolic area, so the estimate should match EF from ground-truth masks.Source: §3.8 · p. 7
- For 2D-mask datasets, EF is an area-based surrogate rather than true volumetric clinical EF, while still measuring how well deformation captures cardiac contraction.The calculation compares a warped end-systolic mask with the end-diastolic mask.Source: §3.8 · pp. 7–8
- Per-sample uncertainty intervals for clinical parameters are derived by Monte Carlo perturbations of the deformation field, scaled by the per-pixel epistemic uncertainty map.The noise is smoothed with a Gaussian kernel to produce spatially correlated fields; the 95% intervals use the distribution's 2.5th and 97.5th percentiles.Source: §3.8 · p. 8
4 Results
Across ACDC and CAMUS, both RUGI variants improve image similarity and anatomical alignment over single-step registration. Ablations distinguish gains from iterative refinement and informative gating, while further evaluations cover pretrained models, computational cost, uncertainty estimates, and cardiac parameters.
4.1 RUGI For Registration
- Both RUGI variants significantly improved all four image-similarity metrics over single-step inference across both datasets.Improvements were consistent across patients and remained significant after correction.Source: Figure 4 · p. 8, §4.1 · p. 8
- Median DSC rose from 73.8% to 83.7% with RUGI-Unc and 84.6% with RUGI-Err in ACDC; in CAMUS it rose from 78.2% to 82.3% and 82.0%.Both methods significantly improved DSC relative to the single-step baseline in both datasets.Source: §4.1 · pp. 8, 10
4.1.1 Ablation Study
- Iterative refinement improved similarity even without gating, but both informative gates outperformed ungated, random, and intensity-based ablations in almost all comparisons.RUGI-Err significantly outperformed all three ablations on both datasets after correction.Source: §4.1.1 · p. 11, Table 1 · p. 11
4.1.2 Application to pre-trained models
- Across VoxelMorph, TransMorph, and CycleMorph, RUGI-Err improved MSE, NCC, and MI without retraining; MSE fell 30.2%, 37.1%, and 27.4%, respectively.SSIM effects varied by architecture, including a significant 14.4% decrease for CycleMorph.Source: §4.1.2 · pp. 11–12, Figure 7 · p. 12
4.1.3 Computational Cost
- Both methods reduced loss monotonically, but required longer inference than the single-step baseline; their relative speed and residual error varied by dataset.RUGI-Unc took 5.2×–10.0× longer than baseline; RUGI-Err took fewer steps in ACDC but more in CAMUS.Source: §4.1.3 · pp. 12–13, Table 2 · p. 12, Figure 8 · p. 13
4.2.1 Validation of Estimated Uncertainty
- The top 20% most uncertain and highest-error pixels overlapped by 51.1% in ACDC and 60.7% in CAMUS, versus 20% expected by chance.At a 1-pixel tolerance, overlap reached 69.1% and 76.0%, respectively; RUGI-Unc maps represent predicted epistemic uncertainty.Source: §4.2.1 · p. 13, Figure 9 · p. 14
4.2.2 Diagnostic parameters
- EF MAE fell from 27.5% to 14.9% with RUGI-Unc and 16.0% with RUGI-Err in ACDC; in CAMUS it fell from 30.9% to 23.0% and 25.2%.All pairwise comparisons against the BNN baseline were significant; RUGI slopes were closer to the identity line.Source: §4.2.2 · p. 14, Figure 10 · p. 15
- RUGI’s Jacobian folding rate remained below 2% across methods and datasets, while per-sample uncertainty intervals were available for RUGI-derived cardiac parameters.RUGI variants showed marginally larger deviations from unit Jacobian than the BNN baseline.Source: §4.2.2 · p. 15, Table 3 · p. 15
5 Discussion
RUGI iteratively refines single-pass deformations, and both uncertainty- and error-gated variants improve registration across cardiac MRI and ultrasound. Ablations distinguish the benefit of iteration from the additional benefit of spatial gating, while revealing trade-offs and limitations.
5 Discussion
- RUGI iteratively refines single-pass deformations, modulating updates to focus effort where predictions are less reliable.The discussion describes suppression of potentially destabilizing updates in already well-registered regions.Source: §5 · pp. 15–16
- Both RUGI variants significantly improve over single-step inference across MRI and ultrasound on all four image-similarity metrics.The discussion attributes the pattern to correcting residual misalignment after a single forward pass.Source: §5 · p. 16
- Even no-gating and intensity-based ablations improve over single-step inference, while selective gating provides further accuracy gains.The random-map ablation causes the largest degradation; all ablations perform worse than both proposed gating methods.Source: §5 · p. 16
- The two gating strategies perform similarly on image similarity; uncertainty-based gating requires training, while error-based gating is simpler and model-agnostic.Uncertainty-based gating is consistently better in ejection-fraction estimation, but conclusions are limited to two datasets.Source: §5 · p. 16
- Iterative inference costs more than a single pass, and uncertainty-based gating requires supervised training tied to base-model quality.Error-based gating instead depends on the image residual being informative; extending 2D experiments to 3D may add computational constraints.Source: §5 · p. 16
- Gains are smaller in CAMUS than ACDC, consistent with ultrasound’s lower signal-to-noise ratio, variable speckle, and out-of-plane motion complicating correspondence.The authors suggest multi-scale refinement, additional inference-time regularization, or anatomical structural priors as possible improvements.Source: §5 · p. 16
6 Conclusion
RUGI refines deformation fields at inference by focusing updates on uncertain regions. It improved registration across cardiac MRI and ultrasound, including when error gating was applied to pretrained models without retraining.
6 Conclusion
- RUGI uses uncertainty-gated deformation-field refinement at inference to adapt updates to spatially heterogeneous difficulty.The conclusion reports statistically significant improvements across cardiac MRI and ultrasound, in intensity-based metrics and anatomical alignment.Source: §6 · pp. 16–17
- Error-gated RUGI reduced MSE by 27–37% across VoxelMorph, TransMorph, and CycleMorph without retraining.The conclusion also reports statistically significant improvements across cardiac MRI and ultrasound datasets.Source: §6 · p. 17
- Improved estimation of diagnostic parameters indicates that registration differences can be meaningful for downstream functional-parameter estimation.The uncertainty maps also provide spatially interpretable estimates of registration reliability that can be propagated to downstream tasks.Source: §6 · p. 17