Spatial regularisation for improved accuracy and interpretability in keypoint-based registration

TL;DR

This work tackles the interpretability bottleneck in unsupervised keypoint-based medical image registration by introducing a principled three-fold spatial regularisation loss. The components $L_{KL}$, $L_{var}$, and $L_{rep}$ shape feature maps into interpretable point spread functions, sharpen landmark precision, and enforce landmark diversity, respectively, within a unified training objective $L_{training} = L_{sim} + \lambda_{KL} L_{KL} + \lambda_{var} L_{var} + \lambda_{rep} L_{rep}$. When applied to EquiTrack and KeyMorph, the regularisation yields accurate, anatomically meaningful keypoints, significantly improving rigid motion tracking in foetal MRI and affine brain MRI registration, and approaching supervised performance while maintaining unsupervised learning advantages. The approach enhances interpretability and reliability of landmark-based registration in clinical imaging, with code available for reproducibility and potential extensions to cross-modality scenarios and dynamic keypoint counts.

Abstract

Unsupervised registration strategies bypass requirements in ground truth transforms or segmentations by optimising similarity metrics between fixed and moved volumes. Among these methods, a recent subclass of approaches based on unsupervised keypoint detection stand out as very promising for interpretability. Specifically, these methods train a network to predict feature maps for fixed and moving images, from which explainable centres of mass are computed to obtain point clouds, that are then aligned in closed-form. However, the features returned by the network often yield spatially diffuse patterns that are hard to interpret, thus undermining the purpose of keypoint-based registration. Here, we propose a three-fold loss to regularise the spatial distribution of the features. First, we use the KL divergence to model features as point spread functions that we interpret as probabilistic keypoints. Then, we sharpen the spatial distributions of these features to increase the precision of the detected landmarks. Finally, we introduce a new repulsive loss across keypoints to encourage spatial diversity. Overall, our loss considerably improves the interpretability of the features, which now correspond to precise and anatomically meaningful landmarks. We demonstrate our three-fold loss in foetal rigid motion tracking and brain MRI affine registration tasks, where it not only outperforms state-of-the-art unsupervised strategies, but also bridges the gap with state-of-the-art supervised methods. Our code is available at https://github.com/BenBillot/spatial_regularisation.

Figures (3)

• Figure 1: (A) Registration by unsupervised keypoint detection. Beyond the classical similarity loss between fixed and moved volumes (leading to features comparable to B,left), we propose a three-fold spatial regularisation, whose effects on the features/keypoints are depicted in (B-D): interpretability ($\mathcal{L}_{KL}$,), precision ($\mathcal{L}_{var}$), and diversity ($\mathcal{L}_{rep}$).
• Figure 2: Representative features and summary of all extracted keypoints obtained from an EPI foetal brain MRI with EquiTrack, trained with/without different regularisations. Keypoints are colour-coded by their distance to the visualised slice. When trained with our loss (green), EquiTrack yields interpretable, precise and diverse keypoints.
• Figure 3: (A) Average Dice scores for ADNI to HCP brain MRI registration. Stars indicate statistical significance with KM-unsup-ours (1% level, Wilcoxon signed-rank test). (B) Example of HCP/ADNI scans with overlaid segmentations, and representative features from KeyMorph trained with different losses. (C) Input ADNI scan affinely regist-ered by all methods and segmentation of the fixed HCP scan. Arrows show large errors.