Improving Generalization of Medical Image Registration Foundation Model

TL;DR

This paper enhances the generalization and robustness of medical image registration foundation models by integrating Sharpness-Aware Minimization, leading to improved cross-dataset performance and stability across diverse clinical scenarios.

Contribution

It introduces the use of Sharpness-Aware Minimization in foundation models to improve their generalization and robustness in medical image registration tasks.

Findings

Significant improvement in cross-dataset registration accuracy.

Enhanced model stability across diverse data distributions.

Better handling of complex clinical scenarios.

Abstract

Deformable registration is a fundamental task in medical image processing, aiming to achieve precise alignment by establishing nonlinear correspondences between images. Traditional methods offer good adaptability and interpretability but are limited by computational efficiency. Although deep learning approaches have significantly improved registration speed and accuracy, they often lack flexibility and generalizability across different datasets and tasks. In recent years, foundation models have emerged as a promising direction, leveraging large and diverse datasets to learn universal features and transformation patterns for image registration, thus demonstrating strong cross-task transferability. However, these models still face challenges in generalization and robustness when encountering novel anatomical structures, varying imaging conditions, or unseen modalities. To address these limitations, this paper incorporates Sharpness-Aware Minimization (SAM) into foundation models to enhance their generalization and robustness in medical image registration. By optimizing the flatness of the loss landscape, SAM improves model stability across diverse data distributions and strengthens its ability to handle complex clinical scenarios. Experimental results show that foundation models integrated with SAM achieve significant improvements in cross-dataset registration performance, offering new insights for the advancement of medical image registration technology. Our code is available at https://github.com/Promise13/fm_sam}{https://github.com/Promise13/fm\_sam.