BrainMorph: A Foundational Keypoint Model for Robust and Flexible Brain MRI Registration

2024-05-22

AI Summary

To address challenges in brain MRI registration—including large initial misalignments, poor scalability, and the need for handling multi-modal, paired, and large-scale population data—this paper introduces KeyMorph, the first foundation model for general-purpose brain registration. Built upon a keypoint-driven architecture, KeyMorph is trained via large-scale self-supervision on over 100,000 3D neuroimaging volumes, jointly optimizing multi-modal alignment and group-wise registration while enabling automatic keypoint detection and disentangled deformation modeling. Its key innovations include on-the-fly generation of rigid, affine, or controllable non-rigid transformations at inference time, ensuring robustness, interpretability, and precise generative control. Extensive experiments demonstrate that KeyMorph consistently outperforms state-of-the-art methods across diverse registration tasks—particularly excelling in scenarios with severe initial misalignment and large-scale (thousand-subject) population registration, where it achieves higher accuracy and significantly faster inference.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multimodal Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• 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

Abstract

We present a keypoint-based foundation model for general purpose brain MRI registration, based on the recently-proposed KeyMorph framework. Our model, called BrainMorph, serves as a tool that supports multi-modal, pairwise, and scalable groupwise registration. BrainMorph is trained on a massive dataset of over 100,000 3D volumes, skull-stripped and non-skull-stripped, from nearly 16,000 unique healthy and diseased subjects. BrainMorph is robust to large misalignments, interpretable via interrogating automatically-extracted keypoints, and enables rapid and controllable generation of many plausible transformations with different alignment types and different degrees of nonlinearity at test-time. We demonstrate the superiority of BrainMorph in solving 3D rigid, affine, and nonlinear registration on a variety of multi-modal brain MRI scans of healthy and diseased subjects, in both the pairwise and groupwise setting. In particular, we show registration accuracy and speeds that surpass current state-of-the-art methods, especially in the context of large initial misalignments and large group settings. All code and models are available at https://github.com/alanqrwang/brainmorph.

Problem

• Develops a keypoint-based model for brain MRI registration
• Supports multi-modal, pairwise, and scalable groupwise registration
• Ensures robustness to large misalignments and rapid transformation generation

Innovation

• Keypoint-based foundation model for MRI registration
• Trained on 100,000 3D volumes dataset
• Supports multi-modal and scalable groupwise registration