Uncertainty Estimation for Pretrained Medical Image Registration Models via Transformation Equivariance

TL;DR

This work tackles the lack of uncertainty estimates in pretrained medical image registration by introducing an inference-time, model-agnostic framework based on transformation equivariance. By perturbing the source image with spatial transforms and analyzing the variability of the composed predictions, the method produces per-voxel uncertainty maps that correlate with true registration error without retraining. The uncertainty is theoretically decomposed into Intrinsic Spread and Bias Jitter, with concrete formulas linking the variance to the underlying error, and is validated across brain, cardiac, and abdominal datasets using three pretrained backbones. Empirically, the approach shows strong to moderate agreement with MC-dropout and meaningful case studies demonstrating its ability to flag anatomically inconsistent regions, enabling safer, risk-aware deployment of pretrained registration models in clinical and research settings.

Abstract

Accurate image registration is essential in many medical imaging applications, yet most deep registration networks provide little indication of when or where their predictions are unreliable. Existing uncertainty estimation approaches, such as Bayesian methods, ensembles, or MC-dropout, typically require architectural modifications or retraining, precluding their applicability to pretrained registration models. We propose an inference-time, model-agnostic uncertainty estimation framework that applies directly to any pretrained registration network. Our approach is grounded in the transformation equivariance property of image registration, which states that the underlying anatomical mapping should remain consistent under spatial perturbations of the input. Experiments across three pretrained registration models and four anatomical structures show that the resulting uncertainty maps consistently correlate with registration error and highlight unreliably aligned regions. This framework turns pretrained registration networks into risk-aware tools at test time, moving medical image registration closer to safe clinical and large-scale research deployment.

Figures (14)

• Figure 1: Transformation equivariance for image registration. The mapping between image A and B should stay consistent regardless of whether we warp image A by $\tau$ or not.
• Figure 2: Overview of the proposed uncertainty estimation framework. We present an inference-time, model-agnostic framework for uncertainty estimation in image registration that requires no model retraining or architectural modifications. Concretely, we (i) perturb the source image coordinate frame $I^{A^{\prime}}=I^A\circ\tau$, (ii) compose the predicted transformations $\tau\circ\hat{\phi}^{A^{\prime}B}$ to represent the mapping between the original images $I^A$ and $I^B$, and (iii) measure their variance $\operatorname{Cov}_\tau[\tau\circ\hat{\phi}^{A^{\prime}B}]$ to produce uncertainty maps. This variance decomposes into two interpretable components: Intrinsic Spread, reflecting the average local variability of the error distribution, and Bias Jitter, capturing the shifts of the error distributions under spatial perturbations.
• Figure 3: Comparison between our method and the MC-dropout on single cases. The correlation is computed between the 3D uncertainty maps.
• Figure 4: Quantitative comparison between our method and the MC-dropout with the IXI dataset.
• Figure 5: Correlation between the estimated uncertainty and the true registration error for a combination of varying ground truth transformation types (GT), perturbation transformation types, and datasets with uniGradICON.

Theorems & Definitions (6)

• Lemma 4.1: Mean and covariance under an arbitrary diffeomorphic perturbation
• Lemma 4.2: Mean and variance under affine perturbation
• Corollary 4.3: Translation perturbations
• proof
• proof
• proof