Fine-Tuning TransMorph with Gradient Correlation for Anatomical Alignment

TL;DR

Unsupervised brain MRI registration often yields deformations that lack biological plausibility. The authors fine-tune a pre-trained TransMorph using the Fisher Adam ($FAdam$) optimizer and introduce a gradient correlation loss ($L_{GC}$) defined via $GC(A,B)=\frac{1}{3}(NCC(\nabla_x A,\nabla_x B)+NCC(\nabla_y A,\nabla_y B)+NCC(\nabla_z A,\nabla_z B))$, with $L_{GC}=1-GC(I_t,I_s\circ\phi)$ and $L_{sim}=L_{IC}+\gamma L_{GC}$, $\gamma=0.5$, $L= L_{sim}+\lambda L_{reg}$, $\lambda=2$. Data from Learn2Reg LUMIR (over $4000$ preprocessed T1 MRIs) are used to evaluate the approach, comparing Baseline, $FAdam$, and $FAdam+GC$. Results show that $FAdam+GC$ achieves slight gains in Dice and $HdDist95$ and a notable reduction in non-diffeomorphic volume (NDV) relative to Baseline, with qualitative evidence of improved boundary alignment. The study demonstrates that gradient correlation improves deformation smoothness and structural consistency in unsupervised intersubject brain MRI registration, and that fine-tuning pre-trained models can enhance performance with limited retraining.

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.

Figures (3)

• Figure 1: Moving, fixed, moved ("Baseline") and moved ("FAdam+GC") images. Arrows indicate some example areas where the smoothness of the deformation was improved.
• Figure 2: Displacement fields and grids of "Baseline" and "FAdam+GC" - as RGB image coding XYZ displacements (left) and as grid (right). The arrows indicate some example areas where the smoothness of the deformation was improved (cfg. figure \ref{['fig:dis_grid']}).
• Figure 3: GC values of "Baseline" and "FAdam+GC" as well as the difference in GC and intensity values'. The arrow indicates an example of an area where an artifact created by the baseline was reduced.