Hierarchical Uncertainty Estimation for Learning-based Registration in Neuroimaging
Authors: Xiaoling Hu
Karthik Gopinath
Peirong Liu
Malte Hoffmann
Koen Van Leemput
Oula Puonti
Juan Eugenio Iglesias
Venue: ICLR 2025 Poster (2025)
Keywords: Image registration
uncertainty estimation
medical image analysis
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose a principled way to propagate uncertainties (epistemic or aleatoric) estimated at the level of spatial location by these methods, to the level of global transformation models, and further to downstream tasks.

Abstract: Over recent years, deep learning based image registration has achieved impressive accuracy in many domains, including medical imaging and, specifically, human neuroimaging with magnetic resonance imaging (MRI). However, the uncertainty estimation associated with these methods has been largely limited to the application of generic techniques (e.g., Monte Carlo dropout) that do not exploit the peculiarities of the problem domain, particularly spatial modeling. Here, we propose a principled way to propagate uncertainties (epistemic or aleatoric) estimated at the level of spatial location by these methods, to the level of global transformation models, and further to downstream tasks.  Specifically, we justify the choice of a Gaussian distribution for the local uncertainty modeling, and then propose a framework where uncertainties spread across hierarchical levels, depending on the choice of transformation model. Experiments on publicly available data sets show that Monte Carlo dropout correlates very poorly with the reference registration error, whereas our uncertainty estimates correlate much better. Crucially, the results also show that uncertainty-aware fitting of transformations improves the registration accuracy of brain MRI scans. Finally, we illustrate how sampling from the posterior distribution of the transformations can be used to propagate uncertainties to downstream neuroimaging tasks. Code is available at: https://github.com/HuXiaoling/Regre4Regis.

Published as a conference paper at ICLR 2025
HIERARCHICAL UNCERTAINTY ESTIMATION FOR LEARNING-BASBII REGISTRATION IN NEUROIMAGING
Xiaoling Hu Karthik Gopinath', Peirong Liu', Malte Hoffmann', Koen Van Leemput
Oula Puonti 1331, Juan Eugenio Iglesias 14,5,3
Massachusetts General Hospital and Harvard Medical School
2 Aalto University
3 Danish Research Centre for Magnetic Resonance, Copenhagen University Hospital
"Hawkes Institute. University College London
5 Computer Science and AI Laboratory, Massachusetts Institute of Technology
ABSTRACT
Over recent years, deep learning based image registration has achieved impres- sive accuracy in many domains, including medical imaging and, specifically, hu- man neuroimaging with magnetic resonance imaging (MRI). However, the uncer- tainty estimation associated with these methods has been largely limited to the application of generic techniques (e.g.. Monte Carlo dropout) that do not exploit the peculiarities of the problem domain, particularly spatial modeling, Here, we propose a principled way to propagate uncertainties (epistemic or aleatoric) es- timated at the level of spatial location by these methods, to the level of global transformation models, and further to downstream tasks. Specifically, we jus- tify the choice of a Gaussian distribution for the local uncertainty modeling, and then propose a framework where uncertainties spread across hicrarchical levels, depending on the choice of transformation model. Experiments publicly avail- able data sets show that Monte Carlo dropout correlates very poorly with the ref- erence registration error, whereas our uncertainty estimates correlate much better. Crucially, the results also show that uncertainty-aware fitting of transformations improves the registration accuracy of brain MRI scans. Finally, illustrate how sampling from the posterior distribution of the transformations can be used to propagate uncertainties downstream neuroimaging tasks. Code is available at: https://qithnne http://hwha....oo///oooooooooooa
INTRODUCTION
Aligning two more images to common coordinate frame, referred to image registration, is one of the fundamental tasks in medical image analysis, especially in human neuroimaging with MRL. fairly consistent across subjects, registration methods have
Since many human brain structures been very successful in this domain. Registration plays vital role in many important applications. One example is measuring temporal change in longitudinal studies by registering scans of the same subject at different time points (Holland et al., [) by using each subject as its own control, confounding effects introduced by inter-subject morphological variability are considerably reduced. Another important application has been image segmentation: prior to deep learning, the state of the art was established by multi-atlas approaches (Iglesias & Sabuneu (015).. based on registering set of labeled cases to the target scan and merging the deformed segmentations. Yet another key application has been spatial normalization
digital atlas (e.g., the ubiquitous MNI template, Fonov et al. 2009), which is at the core of many neuroimaging studies, and which enables analyses (c.g., regression, group comparison) as a precise function of spatial location (Sowell et al., 2003).
Classically, image registration is cast an optimization task, where the aim is to maximize mea- sure of similarity between a pair of images with respect to a (non)-linear transformation - often combined with regularization term that prevents excessively convoluted deformations (Zitova & Flusser 2003). This problem is typically solved with standard numerical optimization methods
Email: Xiaoling Hu (xihu3@ www.ae@@mghhaaaa.ddu... Co-senior authors

---

Published as a conference paper at ICLR 2025
(Nocedal & Wright 1999). In medical imaging, the different components of optimization based registration have been exhaustively studied, including transformation models (Rueckert et al. 1999.
Christensen & Johnson 2 similarity metrics (Pluim et al. 2003) Avants et al.) 2008). and op-
timization approaches (Klein et al. 2007 Glocker et al. 2011 The reader is referred to (Sotiras)
et al. 2013) for a comprehensive survey.
During the last decade, the focus on image registration has shifted from optimization-baase ap- proaches to (deep) learning-based approaches. By sidestepping the numerical optimization task, these algorithms can predict a mapping between two images almost instantaneously. Earlier deep learning (DL) methods, best represented by QuickSilver (Yang et al., 2017). were trained in a su- pervised manner using ground-truit deformation fields obtained with classical optimization-bass registration. These supervised approaches were superseded by unsupervised frameworks that learn to directly minimize the dissimilarity between two input images, using losses similar to those of clas- sical methods- - p combined with segmentation losses for improved alignment of anatomical regions (Balakrishnan et al., 2019: De Vos et al., 2019). By now, many features of classical reg- istration, including diffcomorphisms (Krebs et al., 2019), symmetry (Iglesias 2023), progressive
warping (Lv et al. 2022), inter-modality support (Hoffmann al. 202 have been incorporated into DL registration.
A crucial, but less explored, aspect of image registration is uncertainty estimation. Mcasures of uncertainty can enhance both the reliability and interpretability of the registration process, which crucial aspects in many downstream applications. In high-stakes clinical applications, such as surgery or radiation therapy, uncertainty estimates can be used to highlight regions with potential registration errors, e.g., when registering pre- and intra-operative images (Simpson al. 201 201 or to give more reliable estimates of structure borders, e.g., when estimating margins around or- gans at risk in radiation therapy (Risholm et al., 2011). In neuroimaging studies, registration errors propagate to downstream analysis tasks, such as segmentation, group-level statistics, or functional mapping, and uncertainty estimates can be used weigh or exclude data from unreliable regions to compute e.g., improved spatial statistics (Simpson et al. [b) In the classical registration lit- crature, uncertainty models often rely on formulating the task as a probabilistic model, where the transform is random variable (Simpson et al. 2013 Le Folgoe et al. 2017 Risholm et al.) [013] Le Folgoc et al. 2016; Kybic, 2009; Agn & Van Leemput 2019). and where the uncertainty is de- scribed by the posterior probability distribution of the transform. However, uncertainty estimation remains largely unexplored in the modern DL registration literature. While there are well-established approaches in the DL literature (particularly ensembling and Monte Carlo dropout, Gal & Ghahra- 2015 Lakshminarayanal et al. 2017), they operate at the voxel level - rather than on the whole transform. For this reason, uncertainty estimates in registration are scarce and typically a by-product of model predictions (Dalca et al. 2019b Gong et al. (0222 Krebs et al., 2018).
Here, we propose to integrate uncertainty estimation into DL registration in a principled manner. Specifically, we propose method for uncertainty-awant fitting of transformation models to predic- tions made independently at different locations (typically each voxel), which can directly capi- talize on existing DL uncertainty estimation approaches. In this framework, DL solves a simpler docation-by-locatinn regression task, where a network is trained to predict a deformation vector per location (or, alternatively, triplet of target coordinates Gopinath et al. 2024), along with uncertainty estimates (alcatoric and/or epistemic). We can then fit multiple transformation models to the set of predictions; our methods are general and support, e.g., affine transforms, B-splines (Rueckert et al. 1999), or non-parametric (Avants et al.] (208)) transforms. The uncertainty estimates can then be propagated to the model parameters in closed form, enabling: (i) weighted fit, where uncertain locations contribute less to the fitting; and (ii) uncertainty estimation of the model parameters, e.g., B-spline coefficients. The uncertainty on the model parameters effectively considers dependencies across spatial locations and can be further propagated to downstream tasks, e.g., registration-basee segmentation.
(Corpinates
In sum, our method models uncertainty as it propagates through a hierarchy of levels (network output, transform models parameters, downstream tasks), in a principled way that enables sampling, investigating modes of variation, and computing of error bars. Specifically, the main contributions of this work are:
1. We propose framework for propagating network estimated uncertainties (epistemic and/or aleatoric) to transformation models and further to downstream tasks. Furthermore, our
2

---

Published as a conference paper at ICLR 2025
framework allows fitting of multiple different transformation models, which can be flexibly chosen based on the application without retraining the network.
2. Using an experimental setup based on recently proposed coordinate-regrestne method for atlas registration, we show that aleatoric uncertainty estimates correlate well with registration error but the epistemic uncertainty (estimated with Monte Carlo dropout) does not. Crucially, we also show that incorporating the aleatoric uncertainty into the fitting of the transformation model increases registration accuracy.
DL
2 RELATED WORK
Deep learning based medical image registration. During the last decade, deep learning methods have dominated medical image registration (Cao al., 2017 7 Krebs al. 2017 Rohé et al. 2017
Uzunova et al. 2017 Sokooti et al., 2017, Tian et al. 2022; Balakrishnan et al. 2019 Dalca et al., b).. These methods directly predict deformation field given two images, and can be gen- erally categorized as supervised and unsupervised. Supervised methods train deep neural networks with ground truth deformation fields. For example, the pioneering QuickSilver method (Yang et al.) 2017) uses fields estimated with an accurate, computationally expensive Large Deformation Diffco- morphic Metric Mapping (LDDMM) model. Cao et al. (2017) learns the complex mapping from the input patch pairs their respective deformation field in patch-based manner. In the context of prostate imaging, Krebs et al. (2017) investigates how deep learning could help organ-specific deformable registration, in applications such motion compensation or atlas-based segmentation. Uzunova et al. (2017) seeks to learn highly expressive appearance models from a limited number of training samples.
Requiring ground truth deformation fields has the disadvantage that trying to learn the distribution of such fields in isointense image regions may lead to wasted model capacity and misguiding gradients due to overfitting. In contrast, unsupervised registration usually uses spatial transformer networks (STN) (Jaderberg et al. 2013) to warp moving images to match fixed images, and the model param- eters are trained using the similarity between warped and fixed images. Similarly to classical meth- ods, regularization terms are often used to encourage the smoothness of the predicted displacement fields. Representative unsupervised methods include VoxeIMorph (Balakrishnan et al. 2019), its variational extensions (Dalca et al.] 2019b), and (De Vos et al.) 2019). These unsupervised methods achieve accuracy levels comparable with classical techmiques, albeit with much higher efficiency.
Uncertainty estimation. Uncertainty estimation seeks assess how confident model in its predictions- - which is of great importance in the deployment of models in the real world, particularly in critical applications such as medical imaging. DL models deal with two types of uncertainty: aleatoric and epistemic. The former is input-dependent, c.g., noise in the data, and can be learned during training (Malinin Gales 2018). Epistemic uncertainty on the model weights, e-g.. due to insufficient training data. Principled formalisms such as Bayesian neural networks are possible, but are only practical for smaller models (Kendall & Gal (017). Recently, Monte Carlo (MC) dropout (Gal & Ghahramani] (015) and model ensembles (Lakshminarayanan et al., 2017; Rupprecht al.) 2017) have been proposed as more practical approaches, especially in the context of classification tasks (Abdar al.) 2021: Gawlikowski et al.) 202I). Unfortunately, these approaches also have shortcomings: MC dropout, which randomly turns off a fraction of the neurons, is fast and easy to implement but is known to underestimate uncertainty (Blei et al.] (017). whereas ensembles that estimate uncertainty as variance across multiple networks, or multiple output layers, are more accurate but computationally expensive in training (Lakshminarayanan al. (017).
Uncertainty estimation for image registration. Uncertainty estimation for medical image registra- tion provides a layer of reliability and interpretabillity, and has long been research objective. In the classical literature, methods based on probabilistic modeling have enabled uncertainty estimation via Bayesian inference (Simpson al. [2013] Kybic, 2009: Le Folgoc et al.) 2017] Risholm et al.) 2013 Le Folgoc et al., 2016 Agn & Van Leemput 2019). This is achieved by computing (exactly approximately) the posterior probability distributions of the deformation model parameters. In- stead of Bayesian inference, other methods have used bootstrap sampling as an empirical ensemble method to estimate registration uncertainty (Kybic) 2009).
In the DL era, registration uncertainty is often underutilized, and most existing approaches either rely on direct application of the general uncertainty estimation techniques described above (see for
3

---

Published as a conference paper at ICLR 2025
instance Gong et al. 2022 Smolders et al. 2022 Chen al. 2024a and Chen et al. 2024b for survey). or obtain simplistic uncertainty estimates as a by-product (e.g., the variational inference strategy in Dalca et al. 2019b Sodghi et al. 2019 which is known to underestimate uncertainty). Therefore, these approaches fail to consider the spatial distribution of deformation fields. ZZhan et al. 2024 propose another version of an aleatoric loss for registration, which is parallel the first level of uncertainty, see Section [3.]] in our framework. We propose to propagate the uncertainties further on the transformation models and downstream tasks. Finally, we note that the correlation between uncertainty estimates and registration errors has not been thoroughly investigated (Luo) al. 2019: 2020) possibly due to the scarcity of datasets with labeled pairs of landmarks that
would ideally be used to measure these errors (Luo et al.) 2020).
3
METHODS
Since our framework is applicable to different registration tasks, we first present it in general terms. Specific instantiations of the model, losses, and training approach are presented in Section 3.4
Preliminaries. Given an input image, or input images, I, our goal is to train a network that predicts d d d2, dg] where d = [dy...... dj.v] I" is column vector storing N values for the coordinate direction j. Depending on the application, the values may correspond to target (Ir, y, z) coordinates, target image that are needed to reach voxels an input image, or even key points that can be used for landmark aansed registration (Wang et al. 2023). We are thus seeking to train a neural network F8 with parameters 0, such that: d = FI
displacements from reference voxels
3.1 FIRST LEVEL OF UNCERTAINTY: TARGET REGRESSION
The strength of our method lies in its ability to capitalize on well-established uncertainty estimation methods that operate at the level of the individual outputs di In this context, one can use the tech- niques discussed in Section and Section 2 above to obtain an accurate estimate of the distribution of the field at each point individually. Without loss of generality, we assume this distribution to be Gaussian: Monte Carlo dropout and ensembles both yield samples that are typically summarized into mean and (cr-)vvriann,; whereas aleatoric uncertainty estimation for continuous variables often relies on prediction of Gaussian means and (co-)variances as well (Tanno et al. 2017).
Therefore, and irrespective of the chosen uncertainty modeling approach, we assume throughout the rest of this manuscript the availability of μ = [1,,, f42, [3) and o = 10 22 dz), where stores the predicted mean values for every position and direction and a denotes the corresponding predicted standard deviations. Henceforth, we refer to these as the first level of uncertainty.
We further note that, since uncertainty is considered, overfitting in flat image regions is not issue and one can safely train the network with supervised losses which greatly facilitates learning of aleatoric uncertainty. In this case, ground truth coordinates or displacements can be obtained by registering images to other images (pairwise registration) or MNI (atlas registration) using a slow, accurate, classical method (Yang et al. 2017).
3.2 SECOND LEVEL OF UNCERTAINTY: WEIGHTED FITTING
Given a set of predicted values μ and their standard deviations o,, we can fit a large family of (non- )linear transformations, including models based on basis functions (e.g., B-splines, Rueckert et al. 1999) and non- parameeriiiic approaches (Thirion 1998).
-nwert Parametric Transformations. Let Φ - ( ΦB) be matrix of B basis functions where Φ (6t,1, ...; Øo.N]" !" and Φt,n represents the value of basis function b evaluated at location Finally, let e be B x 1 vector with the coefficients of the basis functions. To fit the model, we minimize the coordinate error weighted by the corresponding precisions, given by the inverse variances. This can be solved in one coordinate direction j at the time using standard weighted least squares. Given a weight matrix W, = diag( 0, 2). the goal is to minimize the weighted squared error E, for each of the three dimensions j 1, 3:
F, = [P,
(025"" Mors]lWWle W,(m, - [,
which has the well-known solution:
4

---

Published as a conference paper at ICLR 2025
w " W Aps.
where of is the mean of the fitted coefficients for coordinate direction J. and Aj (s"W, W 1--d x0 W is the weighted pseudoinverse. Further, since this is linear estimate, we can [
compute the x B covariance matrix of the fitted coefficients:
cf = Aj W * A3.
We that the basis function formulation covers both linear and non-linear transformations, as the former can be seen as special case of the latter with Φ [ Z, 1]. where the first three columns are the coordinates the corresponding N locations in the imput image space and the last one is a column of ones. These linear basis functions can be fitted in isolation (linear registration), prior to nonlinear fitting (sequential linear / nonlinear registration), or together with nonlinear basis functions in single (joint finear/nonlinear registration).
Given and eh for every direction, we can:
Compute the most likely transform as dc
Visualize the variance map of the coefficients diag(c"), e.g., heat map.
Obtain samples of the field as @ W-1g), where g is random vector with zero- mean, unit-variance Gaussians at every entry.
Extract the leading eigenvalues {Aa} and eigenvectors {es} (e.g., with randomized PCA, Rokhlin al. 2010) to visualize the main modes of variation, e.g., $(c + =),e,) with 1,2,3.
k --3333 and
Uncortainty-AAaae -on-Paramctric Transformations. Best represented by the demons algo- rithm (Thirion] 1998), most con-arametric registration methods consist of alternating vector field estimation and field regularization (field smoothing). In our case, the field estimation is given by the network prediction, which is fixed, so iterating is not necessary. Instead, we simply convolve the network output with a smoothing kernel K (typically Gaussian) to obtain the transformation. Since this is a linear operation (independently of the choice of K), the distribution of the smoothed field remains Gaussian, so we can:
Compute the most likely transform, which can be efficiently obtained with convolutions:
K (a- p)/(K o ?)), where @ is the element-wise (Hadamard) product.
Compute the voxel variance map as: (K @ K) x²
Obtain samples of the field as K (1 + W)1g), where is, once again, a random vector with zero-mean, unit-variance Gaussians every entry.
We note that extracting the leading eigenvalues and eigenvectors in this scenario is also possible, but cannot be effectively done with convolutions and requires using the full expression (Equation (T).
Throughout the rest of this paper, we refer to the distribution of the fitted transformation (whether it coefficients e of the smoothed non parameeriii field) as the second level of uncertainty.
3.3 THIRD LEVEL OF UNCERTAINTY: ERROR BARS ON DOWNSTREAM TASKS
We can further propagate the uncertainty of the transformation to downstream tasks. As an example, used registration-basee segmentation. Given that we can draw samples of the spatial transfor- mation, we can propagate multiple versions of an atlas segmentation (see Figure T]. Each sample from the transformation leads to one possible segmentation map, resulting in different versions of segmentation maps, which can be used to estimate a distribution of labels at every spatial location and derive its uncertainty using, e.g., the entropy of this distribution. We call this the third level of uncertainty.
3.4 MODEL INSTANTIATION
We demonstrate our framework using a simple coordinate- coordinatt-regrrssion DL method, called "Registra- tion by Regression" (RbR), for atlas registration that has been recently proposed (Gopinath al.)
5

---

Published as a conference paper at ICLR 2025
)024];
extension pairwise registration is straightforward, by regressing displacements from im- age pairs, rather than atlas coordinates from single images. RbR aims to non- non-linearly align a given input scan with target atlas (specifically the MNI template). In this case, d specifies the target coordinates in the MNI space for every voxel in the input scan. Here we describe the losses used to train the baseline RhR model and the additional losses have used in the training of the extended
model.
Coordinate loss. For the coordinate loss simple & loss between the predicted and ground- truth coordinates at every voxel location:
N
Looord En Ma n=1
mn mn - F& Fl)
where m [m],... mx] I flattened binary foreground mask excluding all non-brain regions and cn, and (T). denote the ground-truth and prediction at location n respectively. An & loss could also be used, but we show in the experiments that this yields worse performance.
Mask loss. In addition to regressing the atlas coordinates, the CNN is also trained to predict the brain mask m the MNI template we use a combination of a standard cross-entropy and Dice losses with empirically chosen relative weights of 0.25 and 0.75.
does not include non-brain structures. For the mask loss
1 E00.5 (M m,) 0.75 Ldice (M, mm )]
Lmakk Σn my mn
in which in is the predicted mask and is ground-truie mask, which is also used in the coordinate loss. In this case the cross eessopy and Dice losses are computed over two classes (foreground and background), where the ground-truie -tunt and predicted masks are one-hot encoded.
Atlas segmentation loss (optional). Given we have prediction à FI),,, we can fit trans- formation model, as outlined in Section 3.2 to map segmentation from the target MNI space to the input. This allows us supervise the network training using a more fine-grained Dice loss at the level of neuroanatomical structures. To this end, we can use the transformation models in Sec- tion 3.2 as a differentiable step in the network and train end-to-end using a segmentation loss. Here, the transformation model contains both the linear and non-linear parts. We use standard Dice loss for the segmentation:
N
Lise Σn M7
(Sp: Sn),
where and are the ground-truth and transformed segmentations respectively, and the Dice loss is now computed over multiple neuroanatomical structures (see Figure D for example segmentation).
Aleatoric uncertainty loss (optional). To model the aleatoric uncertainty, we learn to directly predict the means fi and standard deviations & of the coordinates. We use the hat to emphasize that these are predictions from the network. When modeling the uncertainty with Gaussian distribution aim to minimize the log- ligellkl loss w.r.t. fi and & given the targets d:
Mn [| inol? + log(
Lunner
Σn Mh
62.
Here dn.5 denotes the ground-truth coordinate for direction j and voxel n. In practice, we learn to predict Log(o",,) rather than oi, directly to map the values to real numbers (not only positive real numbers). Another distribution, such the Laplace distribution could also be used, and in practice gives similar performance as shown in the experiments.
Epistemic uncertainty (optional). To model the epistemic uncertainty, we train the network using Monte Carlo dropout. This amounts to randomly switching off a part of the activation functions in more layers. The samples can be generated similarly test time by making repeated predictions while randomly dropping some of the activation functions. These samples are then summarized and â and used in fitting the transformations.

---

Published as a conference paper at ICLR 2025
Coordinates mean
Atles
Fitting
Differentiable Deformation
Legon
Input
UNet
Mask di
GT
Seg. 8
GTs
Figure 1: Overview of the training strategy of the proposed method.
Full loss for different network configurations. The full loss for the proposed approach is:
Tnotal
Loorrd + mmsk Tmask + Aee Leeg + Xuncer Luncer-
When training without uncertainty, simply drop the last term, which models the aleatoric uncer- lainty. This term is also not included when modeling the epistemic uncertainty, i.e., when training with Monte Carlo dropout.
The overall training framework of the network is illustrated in Figure Given an input scan, the net- work regresses both the coordinate means μe and the standard deviations a,, as well foreground mask m. The optional atlas segmentation loss is denoted by the dashed box.
EXPERIMENTS
Training and test data. We use the same training and test data sets Gopinath et al. (2024). The training data consists of high-resolution, isotropic, T1-weighted scans of 897 subjects from the HCP dataset (Van Essen et al.) 2013 and I 1148 subjects from the ADNI (Jack Jr et al. 20088, while the test data set includes the ABIDE (Di Martino et al. 2014) and OASIS3 (LaMontagne et al. 2019) data sets. More details can be found in Appendix A
Implementation details. We use the standard U-net (Ronneberger al. 2015) our backbone. More details can be found in Appendix B
4.1 RESULTS
We first present results every level of uncertainty and qualitatively demonstrate the utility of our approach in downstream task. We then move on to presenting quantitative results on the registration accuracy using Dice scores as quality metric. Finally, we show how the different components of the training affect registration accuracy using ablations. Bolded numbers denote significant differences (t-test, 0.05).
4.2 UNCERTAINTY THROUGHOUT THE HIERARCHY
First-level of uncer- tainty: epistemic and aleatoric certainty. To ensure that the first-level uncertainties are useful and benefit the registration. they should correlate with the voxel-level coordinate prediction enror. The proposed aleatoric uncertainty is simply the re-
10
5
(a) Error (epis.) (b) Var (epis.)
(c) Error (alea.)
(d) Var (alea.)
Figure 2: mated variance in mmm (Var) for the epistemic (epis.) uncertainty (a. b) and for the alcatoric (alca.) uncertainty (c. d).
Coordinate prediction error in millimeters (mm) (Error) and esti-
gressed o², whereas for the epistemic uncertainty we train a separate network using dropout layers
but without the aleatoric uncertainty loss Truncer-

---

Published as a conference paper at ICLR 2025
(a) Sample
(b) Sample
(c) Sample
(d) Sample
(e) Sample
0.04
0.02
10
0.4
0.2
-10
(I) Variance
(g) Diff: 2-11
(h) Diff: 3-2 (i) Diff: 4-+3 (j) Diff: 5-+4
Figure 3: The top row shows samples from the B-spline transformation with their variance and sample-to-sample differences displayed the second row (from left to right). Similarly, the third row shows samples from the Demons transformation with their variance and sample-to-sampe dif- ferences shown the last (from left to right). Samples and their differences. Diff: 2->1 means the difference between sample 2 and sample 1, and the same applies to others.
In Figure 2 we show, for a single subject, the error between the predicted and ground-truth coordi- nates along with the voxel-wise variance estimated with epistemic and aleatoric uncertainties. The aleatoric uncertainty highlights the cortex, which is difficult to register, as region of high variance, whereas the epistemic uncertainty is much more noisy with less structure. We also quantify both the Spearman, which is more robust against outliers, and Pearson correlations between the coordi- nate error and the variance for both uncertainty approaches over all subjects in the validation set. The correlations for the aleatoric uncertainty are (0.601 + 0.019) (Spearman) and (0.476 + 0.021) (Pearson), which are significantly higher than the correlations for the epistemic uncertainty (0.181 + 0.017) (Spearman) and (0.108 + 0.012) (Pearson). Importantly, the aleatoric uncertainty shows a strong correlation with the coordinate prediction error in absolute terms, which allows effective downweighting of mispredicted coordinates when fitting the transformation, shown in the next section. Given its superiority, we only consider the aleatoric uncertainty in the subsequent experi- ments.
Second-level of uncertainty: uncertainty transforms. As outlined in Section 3.2
visualize the uncertainty of the parametric and non-parametric transforms by plotting the diago- nal of their respective covariance matrices. The heat maps generated in such way are shown in Figure gl for the B-spline basis function (10 mm spacing) coefficients and the Demons trans- formation (3 mm kernel). As for the first level, the largest variances for both transforms coincide with the cortex, which is difficult to register as it is highly folded.
(a) Input
(b) B B-spline
(c) Demons
Figure 4: The heat map of uncertainty from B- spline and Demons. Note the B-spline coeffi- cients upsampled to the image size for visu- alization.
We note that the B-spline coefficients have very large variances the background, because it is masked out in the fitting, i.e., there are no tar- get coordinate predictions match. The Spearman and Pearson correlations between variance and the coordinate prediction error are (0.619 + 0.022) and (0.438 + 0.020), respectively, for the Demons transformation. Thus, while the correlations are almost the same for the network pre-
8

---

Published as a conference paper at ICLR 2025
(a) Input
(b) Sample
(c) Sample 2
(d) Sample
(e) Sample 4 (f) Seg entropy
0.5
0.5
(g) Segmentation (h) Sample
(i) Sample 2
(j) Sample (k) Sample 4 (1) Seg entropy
Figure 5: Example samples from the B-spline transformation (top row, (b-e)) and the Demons trans- formation (bottom row (h-k)), along with the input scan (a) and the ground-truth segmentation (g).
The last column (f, I) show the entropy of the sampled segmentation.
dicted aleatoric uncertainty, the estimate becomes more spatially coherent as shown in Figure T.. We do not compute the correlation for the variance of the B-spline coefficients as they are deñned on a lower resolution grid (10 mm spacing).
Third-level uncertainty: uncertainty at downstream tasks. To illustrate how the uncertainty could be used in downstream task we show samples of atlas deformations and associated propa- gated segmentations in Figures 3 55 Both figures highlight the variability of the samples, which would have direct effect any downstream analysis using quantities extracted from the segmen- tations, e.g., regional volume (Desikan et al. 2009). The sample-to-sample differences, along with the variance, are again concentrated on the cortex. This registration-bar based uncertainty, when not accounted for, can decrease the power of downstream statistical analyses, or be mistaken for aging effects if the registration errors correlate with age.
4.3 REGISTRATION ACCURACY
We quantitatively evaluate the effect of the weighted fitting (second-level of uncertainty) using Dice scores computed between the ground-truth segmentation and the transformed atlas segmentation. Table 0 shows the Dice scores for the affine, B-spline (10 mm), and Demons mm kernel) transformation models fitted with and without uncertainty. The fits without uncertainty are done using only the predicted mean, i.e., effectively setting W, to identity matrix. The registration ac- curacy, measured by Dice, improves for both the affine and B-spline transforma- tions when uncertainty used, and stays the same for Demons. We further show qualitative examples of segmentations transformed with and without uncertainty for a single subject in Figure(6] The likely reason for the Demons transformation not benefiting from the uncertainty weighting is that used the Demons transformation in the atlas segmentation loss T.ee-- Thus, the network predicted average coordinates might already be close to optimal and no further weighting needed. Nevertheless, the qualitative examples show differences for all transformation models, including Demons, when uncertainty weighting is used.
Table 1: Registration performance for transformations with and without uncertainties.
Fitting Strategy
ABIDE
OASIS3
0.718 0.038
Affine with uncertainty
0.730 + 0.033
0.682 + 0.055
B. Splppie 10 mm
0.782 1 0.020
0.750
0.034
B-Spline 10 with uncertainty
0.790 + 0.019
0.772
0.030
Demons
0.799 = 0.020
0.783 = 0.029
Demons 3 with uncertainty
0.799 + 0.019
0.783 + 0.028
4.3. ABLATION STUDIES
Seg. loss and its weight. In Table we compare the performance RbR the proposed approach without uncertainty but using the atlas segmentation loss, which was not used in the original RbR model. As expected, the Dice scores of the proposed approach are higher as the model is trained to minimize the Dice loss. We note, however, that incorporating the uncertaint-informe further improves the results as shown above. We further show the effect of changing the loss weight in Table 3
fitting

---

Published as a conference paper at ICLR 2025
(a) Input
(b) Segmentation
(c) Affine
(d) Affine (uncer.)
(c) Demons
(f) Demons (uncer.)
(g) B-spline
(h) B-spline (uncer.)
Figure 6: Example transformations fitted with and without uncertainty. (a) and (b) show the input
scan and segmentation. (c)-(h) show the transformed segmentations.
Table 2: Comparison between RbR and proposed approach with segmentation loss (average Dice).
Method
ABIDE
OASIS3
RbR (Gopinath et al
2024
Proposed
Rbil Gopinath et al. 2024
Proposed
Affine
0.712 = 0.045
0.715 - 0.040
0.007 = 0.057
0.669 + 0.058
Demons 1
0.734 + 0.038
0.764 + 0.025
0.730 + 0.026
0.758 + 0.025
Demons 3
0.739 + 0.038
0.767 0.024
0.728 0.028
0.759 0.030
Demons 5
0.743 + 0.039
0.761 = 0.026
0.722 = 0.031
0.755 + 0.026
B B-Spline 2.5
0,733 +: 0.038
0.766 : 0.024
0.728 : 0.026
0.760 1: 0.032
B. Spline
0.735 1: 0,038
0.764 0.027
0.725 0.027
0.757 0.034
B-Spline 10
0.738 + 0.039
0.758 + 0.023
0,713 - 0.032
0.738 + 0.029
Modeling alcatoric uncertainty. We model the vari- in the uncertainty loss separately for each coordi- nate direction. Alternatively one can use the same vari- ance for each direction, i.c., same uncertainty for z, y, 2- coordinates. This is equivalent to an isotropic Gaussian distribution and would simplify the modeling as the un- certainty would be the same for all directions at each loca- tion. Table f shows the effect of each modeling strategy on the Dice score. The results demonstrate that it is important to model the uncertainty in each direction separately. We also ablate the effect of the distribution in Ta comparing the scores when using Gaussian a Laplacian. The results are very similar, and in light of this, we chose to use the Gaussian distribution as it has nice theoretical properties as mentioned in Section 3.2
Table 3: Ablation study for Aseg (aver- age Dice).
Ascu 0.5
ABIDE
OASIS3
0.734 + 0.083
0.706 = 0.043
0.739 + 0.035
0,734
0.036
0.790 1: 0.019
0.772
0.030
0.771 0.021
0,769
0.019
Table 4: Ablation study for the number of chan- nels to model uncertainty (average Dice).
Table 5: Ablation study for distribution used to model uncertainty (average Dice).
of channels
ABIDE
OASIS3
Distribution
ABIDE
OASIS3
Single channel
0.7 + 0.054
0.687 = 0.048
Gaussian 0.785 + 0.025
0.774 + 0.026
Three channels
0.790 1 0.019
0.772 = 0.030
Laplacian
0.790 0.019
0.7220000
5 CONCLUSION
Uncertainty estimates in DL-based registration approaches are often under-utilized although they provide valuable information about registration accuracy. Here we have proposed principled ap- proach for propagating the location-leve uncertainties (first-level), to fitted transformations (second- level), and finally to downstream analyses (third-level). The results: show that aleatoric uncertainty correlates well with the coordinate prediction while epistemic uncertainty does not; show that incorporating the aleatoric uncertainty in the transformation fitting improves regis- tration accuracy; and illustrate how the generated samples could benefit downstream analyses. In the future, we aim to extend the framework to pairwise registration of any two input images, and further validate the utility of the uncertainty estimates in aging studies and group comparisons (c.g., between healthy controls and dementia patients).
10

---

Published as a conference paper at ICLR 2025
Acknowledgments. We sincerely thank Yael Balbastre and Bruce Fischl for their helpful dis- cussions. This work is primarily supported by NIH grants 1RFIMH123195, 1R01AG070988, IR01EB031 14, IUMIMH130981, IRFIAG080371, IR2INS138995, and ROO HD10155. Oula Puonti is supported by a grant from Lundbeckfonden (grant number R360-2021-395).
Ethics Statement. As the neuroimaging data used in this study is freely available and accessible to anyone after filling out the data usage agreement no ethical approval was needed.
Reproducibility Statement. We provide the necessary experimental details in Section 4 as well as Appendix [A] Appendix B including data preparation, training and test data sets, network archi- tecture, and other implementation details. The code to reproduce the results will be made publicly available. Although we cannot redistribute the data, all data sets are freely available for download after filling out the data usage agreement.
REFERENCES
Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, Rajendra Acharya, et al. A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information Fusion, 2021.
Mikael Agn and Koen Van Leemput. Fast nonparametric mutual-information-boe registration and uncertainty estimation. In Uncertainty for Safe Utilization of Machine Learning in Medical Imaging and Climical Image-Based Procedures: First International Workshop, UNSURE 2019, and 8th International Workshop, CLIP 2019, Held Conjunction with MICCAI 2019, 2019.
Brian Avants, Charles I. Epstein, Murray Grossman, and James Gee. Symmetric diffeomorphic image registration with cross-correlation: evaluating automated labeling of elderly and neurode- generative brain. MedIA, 2008.
Guha Balakrishnan, Amy Zhao, Mert R Sabuneu, John Guttag, and Adrian Dalca. VoxeIMorph: learning framework for deformable medical image registration. TMI, 2019.
Benjamin Billot, Douglas N Greve, Oula Puonti, Axel Thielscher, Koen Van Leemput, Bruce Fischl, Adrian Dalca, Juan Eugenio Iglesias, et al. Synthseg: Segmentation of brain mri scans of any contrast and resolution without retraining, MedlA, 2023.
David M Blei, Alp Kueukelbir, and Jon D McAuliffe.
Variational inference: review for statisti- cians. Journal of the American statistical Association, 112(518):859-98) 877, 2017.
Xiaohuan Cao, Jianhua Yang, Jun Zhang, Dong Nie, Minjeong Kim, Qian Wang, and Dinggang Shen. Deformable image registration based on similarity-steered enn regression. In MICCAI, 2017.
Junyu Chen, Yihao Liu, Shuwen Wei, Zhangxing Bian, Aaron Carass, and Yong Du From registra- tion uncertainty to segmentation uncertainty. arXiv preprint arXiv:2403.05711 2024a.
Junyu Chen, Yihao Liu, Shuwen Wei, Zhangxing Bian, Shalini Subramanian, Aaron Carass, Jerry L. Prince, and Yong Du. A survey on deep learning in medical image registration: new technolo- gies, uncertainty, evaluation metrics, and beyond, 2024b. URL
2307 15615
Gary E Christensen and Hans J Johnson. Consistent image registration. TMI, 2001.
Adrian Dalca, Marianne Rakic, John Guttag, and Mert Sabuncu. Learning conditional deformable templates with convolutional networks. In NewrIPS, 2019a.
Adrian Dalca, Guha Balakrishnan, John Guttag, and Mert Sabuncu. Unsupervised learning of probabilistic diffcomorphic registration for images and surfaces. MedlA, 2019b.
Bob D De Vos, Floris F Berendsen, Max A Viergever, Hessam Sokooti, Marius Staring, and Ivana Bgumm A deep learning framework for unsupervised affine and deformable image registration. MedlA, 2019.
11

---

Published as a conference paper at ICLR 2025
Rahul S. Desikan, Howard J. Cabral, Christopher P. Hess, William P. Dillon, Christine M. Glaston- bury, Michael W. Weiner, Nicholas Schmansky, Douglas N. Greve, David H. Salat, Randy L. Buckner, Bruce Fischl, and Alzheimer's Disease Neuroimaging Initiative. Automated MRI mea- sures identify individuals with mild cognitive impairment and Alzheimer's disease*. Brain, 132(8):2048 2057, 05 2009. ISSN 0006-8950.
doi: 10. 10.1093/brain/avp123. URL https:
Adriana Di Martino et al. The autism brain imaging data exchange: towards a large-scale evaluation of the intrinsic brain architecture in autism. Mol Psychiatry, 2014.
Vladimir S Fonov, Alan C Evans, Robert C McKinstry, C Robert Almli, and DL Collins. Unbiased nonlinear average age-appropriate brain templates from birth to adulthood. Neurolmage, 2009.
Yarin Gal and Zoubin Ghahramani. Bayesian convolutional neural networks with bernoulli approx- imate variational inference. arXiv preprint arXiv: /tt::::5000225555 2015.
Jakob Gawlikowski, Cedrique Rovile Njieutcheu Tassi, Mohsin Ali, Jongseok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, et al. A survey of 2021.
uncertainty in deep neural networks. arXiv preprint
Ben Glocker, Aristeidis Sotiras, Nikos Komodakis, and Nikos Paragios. Deformable medical im- age registration: setting the state of the art with discrete methods. Annual review of biomedical engineering, 13(1):219-244 2011.
Xuan Gong, Luckyson Khaidem, Wentao Zhu, Baochang Zhang, and David Doermann. Uncertainty learning towards unsupervised deformable medical image registration. In WACV, 2022.
Karthik Gopinath, Xiaoling Hu, Malte Hoffmann, Oula Puonti, and Juan Eugenio Iglesias. Reg- istration by regression (rbr): framework for interpretable and flexible atlas registration. arXiv preprint arXiv:2404.16788 2024.
Malte Hoffmann, Benjamin Billot, Douglas N Greve, Juan Eugenio Iglesias, Bruce Fischl, and Adrian V Dalca. SynthMorph: learning contrast-invariant registration without acquired images.
TMI, 2021.
Dominic Holland, Anders Dale, Alzheimer's Disease Neuroimaging Initiative, et al. Nonlinear registration of longitudinal images and measurement of change in regions of interest. MedIA, 2011.
Juan Eugenio Iglesias. A ready-to-use machine learning tool for symmetric multi-modality registra- tion of brain MRI. Scientific Reports, 2023.
Juan Eugenio Iglesias and Mert R Sabuncu. Multi-atlas segmentation of biomedical images: survey. MedIA, 2015.
Clifford R Jack Jr, Matt A Bemsttein, Nick C Fox, Paul Thompson, Gene Alexander, Danielle Har- vey, Bret Borowski, Paula J Britson, Jennifer L. Whitwell, Chadwick Ward, al. The Alzheimer's disease neuroimaging initiative (ADNI): MRI methods. J. Magn Reson Imaging, 2008
Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. In NeurlPS, 2015.
Alex Kendall and Yarin Gal. What uncertainties do we need bayesian deep learning for computer vision? In NeurIPS, 2017.
Stefan Klein, Marius Staring, and Josien PW Pluim. Evaluation of optimization methods for nonrigid medical image registration using mutual information and b-splines. TMI, 2007.
Julian Krebs, Tommaso Mansi, Hervé Delingette, Li Zhang, Florin C Ghesu. Shun Miao, Andreas K Maier, Nicholas Ayache, Rui Liao, and Ali Kamen. Robust non-rigid registration through agent- based action learning. In MICCAI, 2017.
12

---

Published as a conference paper at ICLR 2025
Julian Krebs, Tommaso Mansi, Boris Mailhé, Nicholas Ayache, and Hervé Delingette. Unsuper- vised probabilistic deformation modeling for robust diffeomorphic registration. In Deep Learn- ing in Medical Image Analysis and Multimodal Learning for Clinical Decision Support: 4th International Workshop, DLMIA 2018, and 8th International Workshop, MI,-CDS 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 20, 2018, Proceedings 4, 2018.
Julian Krebs, Hervé Delingette, Boris Mailhé, Nicholas Ayache, and Tommaso Mansi. Learning probabilistic model for diffeomorphic registration. TMI, 2019.
Jan Kybic. Bootstrap resampling for image registration uncertainty estimation without ground truth. TIP, 2009.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In NewrlPS, 2017.
Pamela J LaMontagne et al. Oasis-3: longitudinal neuroimaging, clinical, and cognitive dataset for normal aging and alzheimer disease. MedRxiv, 2019.
Loic Le Folgoc, Herve Delingette, Antonio Criminisi, and Nicholas Ayache. Quantifying registra- tion uncertainty with sparse bayesian modelling. TMI, 2016.
Loic Le Folgoc, Hervé Delingette, Antonio Criminisi, and Nicholas Ayache. Sparse bayesian regis- tration of medical images for self-tuning parameters and spatially adaptive parametrization of displacements. MedIA, 2017.
Jie Luo, Alireza Sedghi, Karteck Popuri, Dana Cobzas, Miaomiao Zhang, Frank Preiswerk, Matthew Toews, Alexandra Golby, Masashi Sugiyama, William M Wells, et al. On the applicability of registration uncertainty. In MICCAI, 2019.
Jie Luo, Sarah Frisken, Duo Wang, Alexandra Golby, Masashi Sugiyama, and William Wells III. Are registration uncertainty and error monotonically associated? In MICCAI, 2020.
Jinxin Lv, Zhiwei Wang, et al. Joint progressive and coarse-to- ceare--to-fine registration of brain MRI via deformation field integration and non-rigid feature fusion. TMI, 2022.
Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. In NeurlPS, 2018.
Mare Modat, Gerard Ridgway, Zeike A Taylor, Manja Lehmann, Josephine Barnes, David Hawkes, Nick C Fox, and Sébastien Ourselin. Fast free-form deformation using graphics pro- cessing units. Comput Methods Programs Biomed, 2010.
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
Josien PW Pluim, JB Antoine Maintz, and Max A Viergever. Mutual nntualliiforrmation based registration of medical images: a survey. TMI, 2003.
Petter Risholm, James Balter, and William M Wells. Estimation of delivered dose in radiotherapy: the influence of registration uncertainty. In MICCAI, 2011.
Petter Risholm, Firdaus Janoos, Isaiah Norton, Alex J Golby, and William M Wells III. Bayesian characterization of uncertainty in intra-subject non-rigid registration. MedIA, 2013.
Marc-Michel Rohé, Manasi Datar, Tobias Heimann, Maxime Sermesant, and Xavier Pennec, Svf- net: learning deformable image registration using shape matching. In MICCAI, 2017.
Torsten Rohlfing, Image similarity and tissue overlaps surrogates for image registration accuracy: Widely used but unreliable. IEEE Transactions on Medical Imaging, 31(2):153-163, 2012. doi: 10. 10.1109/TM122011211999
Vladimir Rokhlin, Arthur Szlam, and Mark Tygert. A randomized algorithm for principal compo- nent analysis. SIAM Journal on Matrix Analysis and Applications, 3103):1100-1124, 2010.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In MICCAI, 2015.
13

---

Published as a conference paper at ICLR 2025
Daniel Rueckert, Luke Sonoda, Carmel Hayes, Derek LG Hill, Martin O Leach, and David J Hawkes. Nonrigid registration using free-form deformations: application to breast MR images. TMI, 1999.
Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In ICCV, 2017.
Alireza Sedghi, Tina Kapur, Jie Luo, Parvin Mousavi, and William M Wells. Probabilistic image registration via deep multi-class classification: characterizing uncertainty. In Uncertainty for Safe Utilization of Machine Learning Medical Imaging and Climical Iagge-Based Procedures: First International Workshop, UNSURE 2019, and 8th International Workshop. CLIP 2019, Held in Conjunction with MICCAI 2019, Shenzhen, China, October 17, 2019, Proceedings 8, 2019.
Amber Simpson, Burton Ma, Randy E Ellis, A James Stewart, and Michael Miga. Uncertainty propagation and analysis of image-guided surgery. In Medical Imaging 2011: Visualization, Image-Guided Guided Procedures, and Modeling. SPIE, 201
Ivor JA Simpson, MarkW Woolrich, Adrian R Groves, and Julia Schnabel. Longitudinal brain mri analysis with uncertain registration. In MICCAI, 201
Ivor JA Simpson. Mark W Woolrich, Manuel Jorge Cardoso, David M Cash, Mare Modat, Julia A Schnabel, and Sebastien Ourselin. A bayesian approach for spatially adaptive regularisation in non-rigid registration. In MICCAI, 2013.
Andreas Smolders, Tony Lomax, Damien Charles Weber, and Francesca Albertini. Deformable im- age registration uncertainty quantification using deep learning for dose accumulation in adaptive proton therapy. International Workshop on Biomedical Image Registration, 2022.
Hessam Sokooti, Bob De Vos, Floris Berendsen, Boudewijn PF Lelieveldt, Ivana Kgum, and Marius Staring. Nonrigid image registration using multi-scale 3d convolutional neural networks. In MICCAI, 2017.
Aristeidis Sotiras, Christos Davatzikos, and Nikos Paragios. Deformable medical image registration: A survey. TMI, 2013.
Elizabeth R Sowell, Bradley S Peterson, Paul M Thompson, Suzanne E Welcome, Amy L Henke- nius, and Arthur W Toga. Mapping cortical change across the human life span. Nature neuro- science, 6(3):309-31 -315, 2003.
Ryutaro Tanno, Daniel E Worrall, Aurobrata Ghosh, Enrico Kaden, Stamatios N Sotiropoulos, An- tonio Criminisi, and Daniel C Alexander. Bayesian image quality transfer with cnns: explor- ing uncertainty in dmri super-resolution. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, Septem- ber 11-13, 2017, Proceedings, Part 20, pp. 611-619. Springer, 2017.
J-P Thirion. Image matching as a diffusion process: an analogy with Maxwell's demons. MedlA, 1998.
Lin Tian, Yueh Z Lee, Raúl San José Estépar, and Mare Niethammer. Liftreg: limited angle 2d/3d deformable registration. In MICCAI, 2022.
Hristina Uzunova, Matthias Wilms, Heinz Handels, and Jan Ehrhardt. Training enns for image registration from few samples with model-based data augmentation. In MICCAI, 2017.
David C Van Essen, Stephen M Smith, Deanna M Barch, Timothy EJ Behrens, Essa Yacoub, Kamil Ugurbil, Wu-Minn HCP Consortium, et al. The WU-minn human connectome project: an overview. Neuroimage, 2013.
Alan Q Wang, M Yu Evan, Adrian V Dalca, and Mert R Sabuncu. A robust and interpretable deep learning framework for multi-modal registration via keypoints. MedIA, 2023.
Xiao Yang, Roland Kwitt, Martin Styner, and Mare Niethammer. Quicksilver: Fast predictive image
registration-a deep learning approach. Neurolmage, 2017.
14

---

Published as a conference paper at ICLR 2025
Xiaoran Zhang, Daniel H. Pak. Shawn S. Ahn, Xiaoxiao Li, Chenyu You, Lawrence H. Staib, Al- bert J. Simusas, Alex Wong, and James S. Duncan. Heteroscedastic uncertainty estimation frame- work for unsupervised registration. In Marius George Linguraru, Qi Dou, Aasa Feragen, Stamatia Giannarou, Ben Glocker, Karim Lekadir, and Julia A. Schnabel (eds.), Medical Image Computing and Computer Assisted Intervention MICCAI 2024 pp. 651-661, Cham, 2024. Springer Nature Switzerland. ISBN 978-3-031-77066-7.
Barbara Zitova and Jan Flusser. Image registration methods: a survey, Image Vision Comput, 2003.
15

---

Published as a conference paper at ICLR 2025
A TRAINING AND TEST DATA
We use the same training and test data sets as Gopinath et (2024). The training data consists of high-reas reshuueso,ut,on, isotropic, Tl -l-weighted scans of 897 subjects from the HCP dataset (Van Essen jet al. 2013) and 148 subjects from the ADNI (Jack et al. 2008). These two data sets provide good mix of older (ADNI) and younger (HCP) subjects capturing a large range of spatial deforma- tion. The training scans were resampled to Imm isotropic resolution, masked, and registered to the ICBM 2009b Nonlinear Symmetric MNI template using NeftyReg (Modat et al. 2010). Specifically, we first ran the block matching algorithm for affine registration (zeg. Reeg.lladin n) with the -noSym option, and subsequently used the fast free-form deformation algorithm (f3d) to compute the non- linear registration. f3d was run in diffeomorphic mode (-vel) and local normalized correlation coefficient
as similarity metrie. The processing time for the whole dataset less than 24h on a 64-core desktop. To obtain the ground-truth anatomical segmentations, we used SynthSeg Billot et al. (2023) which segments the input scan 32 different structures. The SynthSeg seg- mentation is also used to create the brain mask. The total training set consists of the coordinates, segmentations, and brain masks, and split 80/20% between training and validation.
The test data set consists of two public data sets, ABIDE (Di Martino et al. 2014) and OASIS3 [LaMontagne et al.) 2019), both consisting of high-resolution, isotropic, TI scans. Similar to the training set, we have both younger (ABIDE) and older (OASIS3) subjects for testing. We selected the first 100 scans from both data sets for evaluation so that the test data set matches that Gopinath er al.] (2024). The test data is processed exactly the same way as the training data sets, yielding the coordinates, segmentations and brain masks. The ground truth registration from NiftyReg was visually quality controlled both data sets to check for obvious registration errors both data sets. None of the registrations needed to be excluded.
B IMPLEMENTATION DETAILS
We use the standard U-net (Ronneberger et 2015) our backbone. It has four resolution levels with two convolutional layers (comprising 3 x 3 x 3 convolutions and ReLu) followed by 2 x 2 x 2 max pooling (in the encoder) or upconvolution (decoder). The final activation layer is linear, to regress the atlas coordinates in decimeters (which roughly normalizes them from -1 to 1). We empirically set Amast = 0.5, Aseg and Auneer 0.1. The learning rate is 0.01. The parameters are chosen via the validation performances. The input MRI scan(s) I undergo intensity augmentation with blurring, bias fields, and noise, Furthermore, we also use spatial augmentation, both affine and nonlinear, which is applied to MRI scans well as the segmentations and masks m to ensure spatial correspondence.
C MORE ABLATION STUDIES
Choice of coordinate loss. For the coordinate regression, we can either adopt a L1 or a L2 loss as distance measure. L1 is often used because of its higher robustness, but L2 is usually faster to converge when training. The Dice score difference between the losses is shown in Table 6 The L2 loss performs slightly better than the L1 loss.
Table 6: Ablation study for regression loss (av- erage Dice).
Table 7: Ablation study for fitting strategy (av- erage Dice).
Regression loss
ABIDE
OASIS3
Deformation strategy
ABIDE
OASIS3
L1
0.783 = 0.024
0.763 + 0.016
B-spline
0.710 - 0.022
0.700
0.024
L2
0.790
0.019
0.772 = 0.030
Demons
0.790 0.019
0.772 0.030
Choice of transformation during training. For the atlas segmentation loss, we need to transform the atlas segmentation to input space. We compare B-splines (10 mm spacing) to Demons (3 mm kernel) in Table 7 The non-parametric transformation performs better, however we did not test all possible control point spacings for B-splines or smoothing kernels for Demons. It is possible that another parameter combination would yield even better performance.
16

---

Published as a conference paper at ICLR 2025
D WHY DOES THE UNCERTAINTY IN THE HIT TRANSFORM NOT IMPACT THE RESULT FROM THE DEMONS TRANSFORMATION MODEL?
In the original submission, we speculated that there little gain by fitting with uncertainty when using Demons, because the model training used a Demons fit already in the atlas segmentation loss (see Section 4.3), which limits the margin for improvement when you are fitting this same model at test time.
To further validate our assumption, we conducted additional experiments:
1. Use a B-spline transformation during training instead of the Demons. If our speculation about the limited improvement at test time is correct, fitting the B-spline transformation with uncertainty at test time should show only marginal improvement compared to fitting without uncertainty. As shown in Table g this is indeed the case: the Dice scores for the test-time fits with and without uncertainty are almost exactly the same.
Table Registration performance for transformations with and without uncertainty if training with B-spline.
Fitring Strategy
w/o uncertainty
0.710 0.025
0.699 0.023
w/ uncertaitty
0.710 = 0.022
0,700 = 0.024
2. Alternatively, we trained the model using a direct deformation, i.e., interpolating using the predicted coordinates directly without using a transformation model, in the atlas segmen- tation loss. Now, as shown in Table the Demons transformation does benefit from fitting with using the uncertainty estimates, which further supports our initial suspicion that the improvement for the Demons model was marginal because it was used during training.
Table 9: Registration performance for transformations with and without uncertainties if trained with direct deformation.
Fittng Strategy
ABIDE
OASIS3
w/o uncertainty
0.767
0.024
0.759
0.030
w/ uncertainty
0.799 0.019
0.783 0,028
Based on these results, seems that the uncertainty estimates are helpful in generalizing the fitting accuracy across transformation models (beyond the one used during training).
E COMPUTATIONAL COST
In our experiments, training took approximately 1.2h per epoch on an NVIDIA RTX A6000 GPU, with convergence typically requiring 50 epochs. During inference, the average runtime for a single sample (176 x 256 x 256) was approximately 15 minutes to generate all the results, which includes coordinate prediction and, fitting all of the transformation models (affine, B-spline, Demons) with and without uncertainty. During the inference, most of the time is spent on the uncertainty-aware Bspline fitting, in which we compute the inverse of the high dimension matrix (please refer to Section 3.2). We note that the fitting and sampling steps can be further optimized for speed, and that the inference for single model, e.g., Demons, can be done considerably faster.
F LIMITATIONS
Here we have only included two representative non-linear transformation models and restricted the application to registration to fixed atlas space. Due to computational limitations, we did not com- pute the closed-form solution for the B-spline coefficients for spacings < 10mm. Using Dice scores the only proxy for comparing registration accuracy gives only partial understanding of the qual- ity of the registrations (Rohlfing) (2022). Finally, the usefulness of the third-level uncertainties is presented only qualitatively, and not quantifying in relevant application is left for future work.
17


Reviews:
[{"review_id": "YgmihXuCmt", "reviewer": "Reviewer_ofnx", "summary": "Authors present a method to estimate the uncertainty in medical image registration at 3 different stages: (1) on the estimate of the distribution of the deformation field at each voxel (assuming this is Gaussian and gathering a mean and standard deviation per point), (2) on the distribution of the fitted transformation, and (3) on the distribution of possible outcomes on the downstream task. The uncertainty from the first level is used in the transformation fitting step weighting down contributions from less certain pixels. By drawing samples of the fitted transformation, a distribution of results for the same downstream task is generated which is used to estimate the 3rd level of uncertainty. Authors demonstrate their approach in the context of registration-based Brain image segmentation where the given input image is deformed to match the standard MNI atlas and the labels from the atlas are propagated to the deformed image.", "questions": "* The panel labels in Figure 2 are badly formatted - please make these clearer.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* Authors clearly lay out their proposed framework describing the 3 level of uncertainty they they aim to model.\n* The results are presented such that each level of uncertainty is evaluated - this makes it easy for the reader to better understand the Authors important contribution.\n* Authors provide analysis of the estimates of uncertainty, and demonstrate that the aleotoric uncertainty corresponds better with coordinate prediction error at the first level.\n* The figures in the paper are very useful for visualizing the variations that arise from sampling the fitted transform - they demonstrate the fundamental problem with biomedical image registration in that the results from a downstream task are highly dependend on the transformation paramaters, but the uncertainty can be successfully quantified.", "weaknesses": "* The addition of uncertainty in the fit transform did not impact the result from the Demons transformation model. Authors explain that this is likely because the loss from the Demons method is used in the segmentation loss and that leads to the average coordinates landing close to the optimum. This is unclear - please can Authors expand on this and also elaborate on whether they would expect to see improvements if they used an alternative to Demons?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SpyPEmZ5qo", "reviewer": "Reviewer_X1vX", "summary": "This paper introduces a framework to integrate uncertainty estimation into deep learning-based image registration for neuroimaging. By propagating epistemic and aleatoric uncertainties from voxel-level predictions to transformation models and downstream tasks, the framework enhances registration accuracy and reliability. Experiments show that this uncertainty-aware approach can boost brain MRI registration performance.", "questions": "1. Why is uncertainty estimation important in neuroimaging registration? What specific benefits does it bring, and why is it necessary?\n2. For modeling epistemic uncertainty, why did the author choose to use Monte Carlo dropout, and how exactly was it implemented?\n3. Given that the deformation field (non-linear transformation) is a high-dimensional tensor, how do authors verify the accuracy of the ground-truth labels when collecting them? As I understand, NiftyReg is a registration method—if this method is used to generate ground-truth labels, does that imply the proposed method can only perform as well as NiftyReg at best?\n4. The total training loss includes multiple terms marked as optional in the paper. Are these terms truly necessary? Was an ablation study conducted to assess their impact on performance?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper addresses an important problem in neuroimaging, aiming to use uncertainty estimation to enhance the accuracy and reliability of deep learning-based registration.\n2. The figures and explanations are clear and well-organized, which makes complex ideas easier to understand.\n3. The empirical findings are insightful. The authors compare their proposed method with other different methods.", "weaknesses": "1. The presentation could be improved. The background, motivation, and knowledge gap are not clearly explained, making it challenging to follow the paper’s purpose and direction. Although the study is about why uncertainty estimation matters in brain registration, it doesn’t make a strong case for why this is important. There isn’t enough support or reasoning behind this focus, which leaves readers unsure about the value or impact of the work.\n2. The experiments are insufficient. This paper only compares its method to a single approach, RbR. Although RbR was released in 2024, it is still unpublished and available only on arXiv, which makes it less robust as a baseline for comparison. Including more established and published baselines would be essential to provide a solid foundation for evaluating the effectiveness of the proposed approach.\n3. From the experimental results (e.g., Table 1), the introduction of uncertainty estimation does not appear to provide a statistically significant improvement to the model. This makes it difficult to assess its effectiveness and raises questions about the practical value of incorporating uncertainty estimation into the approach.\n4. Most of the latest work in brain registration focuses on unsupervised learning, largely due to the high cost of collecting accurate transformation ground-truth for high-dimensional images (e.g., 3D MRI). However, this paper still stay on a supervised learning approach, making it seem less aligned with current trends and potentially less practical for real-world applications where labeled data is limited.\n5. The evaluation criteria lack validity. This paper collects ground-truth transformations using NiftyReg, a tool introduced around 15 years ago. Since then, more state-of-the-art methods have been shown to outperform NiftyReg, indicating that the ground-truth it provides may be inaccurate. This undermines the reliability of the results presented in the paper and raises concerns about the credibility of its findings.\n6. The novelty of this work is limited. The main contribution appears to be the addition of an uncertainty loss term, while the other loss terms and network architecture are based on existing work. This incremental improvement does not significantly advance the field, as much of the framework relies on previously established methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oIYTwS69zr", "reviewer": "Reviewer_Eicy", "summary": "In this paper, the authors propose a registration framework that includes uncertainty estimation at three different levels (network output, transform models' parameters and downstream tasks). They show that uncertainty estimates correlate with registration errors and that estimating uncertainty can improve registration performance. The approach is applied to brain MRI registration.", "questions": "- In Figure 2, what does the error correspond to exactly?\n- Could the authors comment on the run time cost of their approach or other considerations regarding its practical application?\n\n===\n- The authors already mention many related works but they could also consider:\n    - Chen J et al. \"A survey on deep learning in medical image registration: New technologies, uncertainty, evaluation metrics, and beyond.\" arXiv preprint arXiv:2307.15615 (2023).\n    - Zhang X et al. \"Heteroscedastic Uncertainty Estimation Framework for Unsupervised Registration.\" International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2024.\nAnd comment especially on the second one.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The framework proposed is flexible and can be applied to a broad range of registration techniques.\n- The authors performed diverse experiments to validate their approach.\n- The results are mostly convincing.\n- The paper is well written and related works well introduced.", "weaknesses": "- The ground truths considered originate from automatic registration and segmentation approaches but it is not mentioned whether quality control was performed and so to which extent they can be relied on.\n- The experiments regarding the downstream task are limited to qualitative results and the fact that the segmentations are not overlaid on the T1w image only allows assessing differences between the segmentation maps but not whether one matches better the anatomy than another one.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)