UniReg: Foundation Model for Controllable Medical Image Registration

2025-03-17

AI Summary

Current learning-based medical image registration methods suffer from poor generalizability, requiring separate model training for distinct clinical scenarios—e.g., inter-subject vs. intra-subject alignment or multi-organ registration—resulting in high computational cost and deployment complexity. To address this, we propose the first controllable foundation model for medical image registration, introducing a conditional deformation field learning paradigm that jointly encodes anatomical priors, registration-type constraints, and instance-specific features, enabling unified multi-scenario registration with a single model. Our architecture integrates anatomy-guided encoding, registration-type-adaptive modulation, and instance feature fusion, eliminating redundant multi-network designs. Evaluated across 90 anatomical structures, the method achieves state-of-the-art performance while reducing training iterations by approximately 50%, significantly lowering computational overhead.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Multi-modal Vision
• Cognitive Modeling & Cognitive Systems: Other Foundations of Cognitive Modeling & Systems

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Learning-based medical image registration has achieved performance parity with conventional methods while demonstrating a substantial advantage in computational efficiency. However, learning-based registration approaches lack generalizability across diverse clinical scenarios, requiring the laborious development of multiple isolated networks for specific registration tasks, e.g., inter-/intra-subject registration or organ-specific alignment. % To overcome this limitation, we propose extbf{UniReg}, the first interactive foundation model for medical image registration, which combines the precision advantages of task-specific learning methods with the generalization of traditional optimization methods. Our key innovation is a unified framework for diverse registration scenarios, achieved through a conditional deformation field estimation within a unified registration model. This is realized through a dynamic learning paradigm that explicitly encodes: (1) anatomical structure priors, (2) registration type constraints (inter/intra-subject), and (3) instance-specific features, enabling the generation of scenario-optimal deformation fields. % Through comprehensive experiments encompassing $90$ anatomical structures at different body regions, our UniReg model demonstrates comparable performance with contemporary state-of-the-art methodologies while achieving ~50% reduction in required training iterations relative to the conventional learning-based paradigm. This optimization contributes to a significant reduction in computational resources, such as training time. Code and model will be available.

Problem

• Lack of generalizability in learning-based medical image registration
• Need for multiple isolated networks for specific tasks
• High computational resource requirements for training

Innovation

• Unified framework for diverse registration scenarios
• Conditional deformation field estimation model
• Dynamic learning encodes anatomical and instance features