Fine-Tuning TransMorph with Gradient Correlation for Anatomical Alignment

2024-12-30

AI Summary

To address issues of irregular deformations, anatomical distortion, and unstable convergence in unsupervised brain MRI registration, this work proposes an enhanced framework built upon the pre-trained TransMorph model. Specifically, we introduce Gradient Correlation—a novel unsupervised similarity metric—into the loss function for the first time and fine-tune the model using the FAdam optimizer. This integration strengthens tissue boundary preservation and enforces anatomically plausible deformation fields, thereby significantly improving registration stability. Quantitative experiments demonstrate that, while Dice score and 95th-percentile Hausdorff distance (HD95) improve marginally, normalized deformation energy (NDV) is substantially reduced, indicating smoother, more regular displacement fields. Qualitatively, registered images exhibit sharper brain tissue boundaries and markedly reduced structural distortions across subjects, effectively mitigating anatomical inconsistency in inter-subject registration.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Computer Vision: Adversarial Attacks & Robustness
• Search and Optimization: Non-convex Optimization

Application Category

• User Modeling, Personalization and Recommendation: Accountability, Transparency, and Ethics for personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

Unsupervised deep learning is a promising method in brain MRI registration to reduce the reliance on anatomical labels, while still achieving anatomically accurate transformations. For the Learn2Reg2024 LUMIR challenge, we propose fine-tuning of the pre-trained TransMorph model to improve the convergence stability as well as the deformation smoothness. The former is achieved through the FAdam optimizer, and consistency in structural changes is incorporated through the addition of gradient correlation in the similarity measure, improving anatomical alignment. The results show slight improvements in the Dice and HdDist95 scores, and a notable reduction in the NDV compared to the baseline TransMorph model. These are also confirmed by inspecting the boundaries of the tissue. Our proposed method highlights the effectiveness of including Gradient Correlation to achieve smoother and structurally consistent deformations for interpatient brain MRI registration.

Problem

• Unsupervised Deep Learning
• MRI Image Matching
• Alignment and Boundary Preservation

Innovation

• Unsupervised Deep Learning
• Gradient Correlation Technique
• Improved TransMorph Model