Recurrent Inference Machine for Medical Image Registration

2024-06-19

arXiv.org

AI Summary

Deep learning-based medical image registration suffers from heavy reliance on large-scale annotated datasets, high computational cost during inference, and limited generalizability across domains. To address these challenges, this paper proposes the Recurrent Inference Meta-Learning framework (RIIR). RIIR pioneers the integration of a latent-state-driven recurrent inference mechanism into registration meta-learning, formulated via differentiable optimization that jointly incorporates explicit gradient inputs and implicit regularization—enabling gradient-guided iterative updates of latent states. This design ensures stable convergence and enhanced generalization under extreme data scarcity. Evaluated on multi-modal brain MRI and quantitative cardiac MRI registration tasks, RIIR achieves state-of-the-art performance using only 5% of the standard training data, significantly outperforming mainstream deep registration methods. Results demonstrate RIIR’s superior accuracy, exceptional data efficiency, and robust cross-domain generalizability.

Technology Category

• Machine Learning: Multi-instance/Multi-view Learning
• Computer Vision: Multi-modal Vision
• Search and Optimization: Metareasoning and Metaheuristics

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Image registration is essential for medical image applications where alignment of voxels across multiple images is needed for qualitative or quantitative analysis. With recent advancements in deep neural networks and parallel computing, deep learning-based medical image registration methods become competitive with their flexible modelling and fast inference capabilities. However, compared to traditional optimization-based registration methods, the speed advantage may come at the cost of registration performance at inference time. Besides, deep neural networks ideally demand large training datasets while optimization-based methods are training-free. To improve registration accuracy and data efficiency, we propose a novel image registration method, termed Recurrent Inference Image Registration (RIIR) network. RIIR is formulated as a meta-learning solver to the registration problem in an iterative manner. RIIR addresses the accuracy and data efficiency issues, by learning the update rule of optimization, with implicit regularization combined with explicit gradient input. We evaluated RIIR extensively on brain MRI and quantitative cardiac MRI datasets, in terms of both registration accuracy and training data efficiency. Our experiments showed that RIIR outperformed a range of deep learning-based methods, even with only $5%$ of the training data, demonstrating high data efficiency. Key findings from our ablation studies highlighted the important added value of the hidden states introduced in the recurrent inference framework for meta-learning. Our proposed RIIR offers a highly data-efficient framework for deep learning-based medical image registration.

Problem

• Improves medical image registration accuracy and data efficiency
• Learns optimization update rules with implicit and explicit regularization
• Outperforms deep learning methods with minimal training data

Innovation

• Meta-learning solver for iterative registration
• Combines implicit and explicit gradient regularization
• Recurrent framework enhances data efficiency