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

2025-03-06

AI Summary

To address feature map spatial dispersion and poor anatomical interpretability in unsupervised keypoint registration, this paper proposes a triple spatial regularization framework: (1) probabilistic keypoint modeling via KL-divergence minimization, explicitly enforcing keypoint distributions to conform to anatomical priors; (2) a feature sharpening loss that enhances response locality and improves localization accuracy; and (3) a repulsion loss that enforces spatial uniformity among keypoints, ensuring geometric plausibility. Integrated into a generic unsupervised registration pipeline, the method enables end-to-end optimization jointly with closed-form point-cloud alignment. Evaluated on fetal rigid-motion tracking and brain MRI affine registration tasks, it outperforms state-of-the-art unsupervised approaches and approaches supervised-method performance, while significantly improving anatomical consistency and interpretability of registration results.

Technology Category

• Machine Learning: Unsupervised & Self-Supervised Learning
• Computer Vision: Motion & Tracking
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Improves interpretability of keypoint-based registration methods.
• Regularizes spatial distribution of features for precise landmark detection.
• Bridges performance gap between unsupervised and supervised registration methods.

Innovation

• KL divergence for probabilistic keypoints modeling
• Sharpening spatial distributions for precise landmarks
• Repulsive loss to encourage spatial diversity