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

TL;DR

KeyMorph is a deep learning framework for multi-modal image registration that uses automatically detected keypoints and a differentiable closed-form solution to improve robustness, interpretability, and handling of large misalignments.

Contribution

The paper introduces KeyMorph, a novel registration method that learns keypoints end-to-end without ground-truth, and leverages a differentiable closed-form for optimal transformation estimation.

Findings

Outperforms state-of-the-art methods in 3D brain MRI registration.

Provides interpretable keypoints that reveal image alignment drivers.

Achieves robustness to large displacements and symmetry considerations.

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 https://github.com/alanqrwang/keymorph.