Learning Diffeomorphism for Image Registration with Time-Continuous Networks using Semigroup Regularization

TL;DR

This work tackles the challenge of learning diffeomorphic image registration by treating the deformation as a time-continuous flow governed by an ODE. It introduces SGDIR, which uses a time-embedded UNet to model φ_t and employs a simple semigroup regularization to enforce the flow map’s composition property, ensuring diffeomorphisms and their inverses without traditional integration schemes or extra regularizers. The approach delivers state-of-the-art performance on four brain MRI datasets in terms of Dice and topology preservation, while also offering competitive results against non-diffeomorphic methods and efficient inference. The framework provides a general, continuous-time perspective for problems that can be cast as autonomous ODEs, with practical impact for robust, topology-preserving medical image registration.

Abstract

Diffeomorphic image registration (DIR) is a fundamental task in 3D medical image analysis that seeks topology-preserving deformations between image pairs. To ensure diffeomorphism, a common approach is to model the deformation field as the flow map solution of a differential equation, which is solved using efficient schemes such as scaling and squaring along with multiple smoothness regularization terms. In this paper, we propose a novel learning-based approach for diffeomorphic 3D image registration that models diffeomorphisms in a continuous-time framework using only a single regularization term, without requiring additional integration. We exploit the semigroup property-a fundamental characteristic of flow maps-as the sole form of regularization, ensuring temporally continuous diffeomorphic flows between image pairs. Leveraging this property, we prove that our formulation directly learns the flow map solution of an ODE, ensuring continuous inverse and cycle consistencies without explicit enforcement, while eliminating additional integration schemes and regularization terms. To achieve time-continuous diffeomorphisms, we employ time-embedded UNets, an architecture commonly used in diffusion models. Our results demonstrate that modeling diffeomorphism continuously in time improves registration performance. Experimental results on four public datasets demonstrate the superiority of our model over state-of-the-art diffeomorphic methods. Additionally, comparison to several recent non-diffeomorphic deformable image registration methods shows that our method achieves competitive Dice scores while significantly improving topology preservation.

Figures (11)

• Figure 1: Schematic of SGDIR. As illustrated in the figure, there are no additional integration or extra loss functions to make the whole training procedure simple. During inference, we can query the model for any desired $t$, e.g., $t=1$ for the final deformation.
• Figure 2: Qualitative comparison of SGDIR with other methods on the OASIS dataset. (visually looks better in color and zoomed in.)
• Figure 3: The visual comparison of the continuous deformations of DiffuseMorph and SGDIR for OASIS dataset.
• Figure 4: The percentage of negative Jacobian determinants through time. The upper figure compares the effect of $\lambda$ on the $|det(J_\phi)<0|$ and the bottom s figure focuses on the best configuration model. The shaded area in the bottom figure indicates the variance across the test set.
• Figure 5: The effect of number of intermediate time steps on the smoothness (percentage of negative Jacobian determinants) and the Dice score.

Theorems & Definitions (10)

• Proposition 1
• proof
• Lemma 1
• proof
• Corollary 1
• proof
• Remark 1
• Lemma 2
• proof
• proof