Improving Generalization of Medical Image Registration Foundation Model

TL;DR

This work tackles cross-dataset generalization in deformable medical image registration by integrating Sharpness-Aware Minimization (SAM) into a foundation-model-based registration framework. The approach uses a multi-scale, Down/TS-based registration network built on GradICON/UNet components and optimizes a SAM objective to flatten the loss landscape, promoting stability across diverse data distributions. Across multi-domain datasets, the SAM-enhanced foundation model delivers superior Dice scores and strongly reduces non-positive Jacobian determinants, indicating both accuracy and diffeomorphic-like regularity, with evidence from ablations and loss-landscape visualizations. The results suggest practical potential for robust, cross-domain medical image registration, with open-source code to facilitate adoption and further development.

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.

Figures (6)

• Figure 1: The illustration shows the combination steps to construct our registration network through downsampling (Down) and the two-step (TS) operator.
• Figure 2: Schematic of the SAM parameter update.
• Figure 3: The boxplot shows the Dice scores of Voxelmorph, Elastix, DiffuseMorph, TransMorph, and our proposed method on 3 anatomical structures in the ACDC dataset.
• Figure 4: An example of axial slices for cardiac MR image registration, displaying the results of all comparison methods. The rows are the target fixed image, the transformed moved image after registration (moved),the registration error map (white indicates zero error, red indicates positive error, blue indicates negative error), the registration transformation, the registered label map (green represents the right ventricle, dark blue represents the myocardium, and light blue represents the left ventricle)
• Figure 5: The boxplot shows the Dice score results of our method and uniGradICON on the ACDC, HippocampusMR, AbdomenMRCT, LungCT, and SLIVER datasets.