A Robust and Interpretable Deep Learning Framework for Multi-modal Registration via Keypoints

Alan Q. Wang,Evan M. Yu,Adrian V. Dalca,Mert R. Sabuncu

2023-09-01

Abstract:We present KeyMorph, a deep learning-based image registration framework that relies on automatically detecting corresponding keypoints. State-of-the-art deep learning methods for registration often are not robust to large misalignments, are not interpretable, and do not incorporate the symmetries of the problem. In addition, most models produce only a single prediction at test-time. Our core insight which addresses these shortcomings is that corresponding keypoints between images can be used to obtain the optimal transformation via a differentiable closed-form expression. We use this observation to drive the end-to-end learning of keypoints tailored for the registration task, and without knowledge of ground-truth keypoints. This framework not only leads to substantially more robust registration but also yields better interpretability, since the keypoints reveal which parts of the image are driving the final alignment. Moreover, KeyMorph can be designed to be equivariant under image translations and/or symmetric with respect to the input image ordering. Finally, we show how multiple deformation fields can be computed efficiently and in closed-form at test time corresponding to different transformation variants. We demonstrate the proposed framework in solving 3D affine and spline-based registration of multi-modal brain MRI scans. In particular, we show registration accuracy that surpasses current state-of-the-art methods, especially in the context of large displacements. Our code is available at <a class="link-external link-https" href="https://github.com/alanqrwang/keymorph" rel="external noopener nofollow">this https URL</a>.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper attempts to address the issue that existing deep learning methods in medical image registration tasks are not robust enough when dealing with large displacements and lack interpretability, failing to effectively utilize the symmetry and equivariance in the problem. Additionally, most models can only produce a single prediction during testing.

Specifically, the paper proposes a deep learning framework based on keypoint detection, called KeyMorph, which aims to achieve multimodal image registration by automatically detecting corresponding keypoints. This framework not only improves the robustness of the registration but also enhances the interpretability of the model, as the keypoints can reveal which parts of the image drive the final alignment. Furthermore, KeyMorph can be designed to have equivariance and symmetry with respect to image translation and the order of input images. Finally, the paper demonstrates how to efficiently compute multiple deformation fields during testing, corresponding to different transformation variants. Through experiments on affine and spline registration of 3D multimodal brain MRI scans, the framework is shown to surpass the current state-of-the-art methods in registration accuracy under large displacements.