Introduction

Inter-modality deformable image registration is crucial to several clinical and research workflows, including motion correction, spatial normalization to templates for segmentation and standardized statistics, and intra-operative registration of portable imaging to clinical scans. While several works have tackled inter-modality image registration^1,2,3,4, they perform considerably worse than methods working in the intra-modality setting due to the difficulty of reliably estimating anatomical similarity across imaging domains.

Most commonly used inter-modality MRI similarity functions include information-theoretic approaches working on intensity histograms (e.g. mutual information and its local extensions^1,2]), approaches focusing on edge alignment (e.g. normalized gradient fields³), and those that build local descriptors invariant to imaging domains (e.g. modality invariant neighborhood descriptors⁴). However, due to their hand-crafted nature, these functions typically require significant domain expertise, non-trivial tuning, and may not consistently generalize outside of the domain-pair that they were originally proposed for.

In this work, we propose a fully unsupervised and data-driven deep network approach to learning an inter-domain similarity function that is sensitive to deformations across multiple scales and used for deformable image registration. Once registered, our approach ensures that corresponding locations have high mutual information in a deep data-driven representation space. Our approach extends recent developments in Patch-based Noise Contrastive Estimation⁵ (PatchNCE) to inter-subject registration by using a pretrained feature space. The proposed method is demonstrated to yield improved registration with higher anatomical overlap as compared to several previous baseline similarity functions.

Materials and Methods

Imaging: Processed 3T T1w and T2w scans and anatomical segmentations are acquired from the S1200 HCP⁶ data release and are downsampled to 2 x 2 x 2 mm³ resolution for rapid prototyping. We use 800/50/191 subjects as training, validation, and testing data for inter-subject registration, respectively.

Training: The overall training pipeline is illustrated in Fig. 1. We align images by maximizing their deep feature similarity and train in two stages.

In stage 1, we train two domain-specific autoencoders with a joint $$$L_1$$$ + Local Normalized Cross Correlation reconstruction objective on their respective training sets. Once trained, their encoders are frozen and used as domain-specific multi-scale feature extractors in stage 2.

In stage 2, a deterministic VoxelMorph registration network⁷ takes a pair of images $$$x_{T1}$$$ and $$$x_{T2}$$$ to be registered as input and yields a stationary velocity field (SVF) $$$v$$$ between $$$x_{T1,T2}$$$ which can be efficiently integrated to obtain a dense displacement field.

Once warped, features are extracted from the (target image, moved image)-pairs by their frozen domain-specific encoders from stage 1. These features are then projected onto a representation space by trainable multilayer perceptrons (as in SimCLR⁸) where the similarity between multiscale features at corresponding locations from the target and moved images can be maximized (as in PatchNCE⁵). We further use a global mutual information loss to maintain global alignment for a final registration loss of (0.1 PatchNCE + 0.9 MI), abbreviated to (NCE + MI) in the remainder of this abstract. To ensure smooth deformations, we employ a diffusion regularizer $$$\|v\|_{2}^2$$$ on $$$v$$$.

Importantly, as all methods are sensitive to regularization strength⁹, we use a hypernetwork conditioned on diffusion regularization strength to sweep over all possible regularization values as in Mok, et al.¹⁰ As a result, all hypernetwork registration models were trained with $$$(1 - \lambda)L_{reg}(x_{T1}, \phi \circ x_{T2}) + \lambda \|v\|_{2}^2$$$ as a loss function where $$$\lambda$$$ is a hyperparameter uniformly sampled from [0, 1] during training and $$$L_{reg}$$$ represents the various similarity functions to be benchmarked.

Evaluation: We benchmark all methods using the commonly used inter-subject registration task^7,9,10 and report Dice overlap between the segmentation of the target image and the moved segmentation of the moving image, indicating registration correctness; and the percentage of folding voxels, indicating deformation quality and invertibility. The proposed loss function was compared against MI¹, Local MI², NGF³, and MIND⁴ as baselines while maintaining the same network architecture. At testing time, we sweep over the range of $$$\lambda$$$ in increments of 0.1 thus generating several registration networks for each benchmarked loss function.

Results

Quantitatively, on a held-out test set of 191 subjects, (Fig 2, left) we see that our method (NCE+MI) with $$$\lambda$$$ = 0.6 - 0.8 outperforms all other losses over all of their respective $$$\lambda$$$’s in terms of anatomical/Dice overlap. Further, (Fig. 2, right) our method (NCE+MI) achieves high registration accuracy while maintaining a comparably negligible percentage of folding voxels to other methods in the $$$\lambda$$$ > 0.6 range, simultaneously indicating high registration accuracy alongside smooth and diffeomorphic deformations. In Fig. 3, we see that the proposed loss function yields more accurate registration perceptually. In Fig. 4, we show that the hypernetwork model generates the expected warped images and deformation fields, with low $$$\lambda$$$’s yielding strong deformations and high $$$\lambda$$$’s yielding highly regular deformations.

Discussion

We present an unsupervised contrastive learning framework for multi-modality diffeomorphic registration. Our results demonstrate improved registration accuracy on a large public database while maintaining smooth and invertible deformations. Experimentally, we validate our claims by benchmarking against several previous hand-crafted loss functions and automatically tune deformation regularity for optimal performance for all methods considered by using hypernetworks. Future work will include extensions of our methodology to other use cases such as high-field to low-field MRI registration and intra-operative multi-modality registration.

Acknowledgements

No acknowledgement found.