$\texttt{NePhi}$: Neural Deformation Fields for Approximately Diffeomorphic Medical Image Registration

TL;DR

NePhi addresses the memory and speed bottlenecks of voxel based medical image registration by introducing an implicit neural deformation field that yields approximately diffeomorphic transformations. It uses dual MLPs for forward and backward deformations, jointly conditioned by a shared latent code $z=[z_g,z_l]$, and enforces approximate invertibility with a gradient inverse consistency regularizer $oxed{\mathcal{L}_{reg}}$. The approach achieves competitive accuracy to voxel based methods in single resolution tasks and substantially lowers memory usage in multi-resolution setups, while enabling fast inference through latent code conditioning and the option for instance optimization. Experiments on 2D synthetic data, COPDGene lung CT, and HCP brain MRI demonstrate strong deformation regularity, scalable memory performance, and practical applicability to high resolution registrations.

Abstract

This work proposes NePhi, a generalizable neural deformation model which results in approximately diffeomorphic transformations. In contrast to the predominant voxel-based transformation fields used in learning-based registration approaches, NePhi represents deformations functionally, leading to great flexibility within the design space of memory consumption during training and inference, inference time, registration accuracy, as well as transformation regularity. Specifically, NePhi 1) requires less memory compared to voxel-based learning approaches, 2) improves inference speed by predicting latent codes, compared to current existing neural deformation based registration approaches that \emph{only} rely on optimization, 3) improves accuracy via instance optimization, and 4) shows excellent deformation regularity which is highly desirable for medical image registration. We demonstrate the performance of NePhi on a 2D synthetic dataset as well as for real 3D medical image datasets (e.g., lungs and brains). Our results show that NePhi can match the accuracy of voxel-based representations in a single-resolution registration setting. For multi-resolution registration, our method matches the accuracy of current SOTA learning-based registration approaches with instance optimization while reducing memory requirements by a factor of five. Our code is available at https://github.com/uncbiag/NePhi.

Figures (11)

• Figure 1: Comparison between NePhi and competing methods based on four design considerations: registration accuracy, regularity of the predicted deformation field, peak memory consumption during training, and inference run-time. NePhi performs slightly worse for pure prediction but on-par with GradICON when combined with instance optimization. In both cases, NePhi shows better deformation regularity and memory consumption compared to other methods. Values are normalized to [0,1] for visualization. Refer to Supp. \ref{['sec:lung_registration_continue']} for details.
• Figure 2: The framework of NePhi. Top: the design of the generalizable neural deformation fields with gradient inverse consistency regularizer (Sec. \ref{['sec:method_neural_deformation_fields']} and Sec. \ref{['sec:method_hybrid_condition']}). Bottom: the overall framework of how to train neural networks to predict NePhi (Sec. \ref{['sec:method_training']}).
• Figure 3: Comparison of peak memory consumption between voxel-representated deformations and NePhi during training and instance optimization (test-time optimization) for 3D images. Details regarding the experiment setting can be found in Sec. \ref{['sec:exp_memory_usage_varying_resolution']}. We observe that memory consumption for voxel-based approaches scales poorly whereas NePhi requires much less memory during training and testing thereby opening opportunities for high-resolution image registration. Such high-resolution capabilities are expected to become increasingly important for example for the registration of very large volumetric miscroscopy images.
• Figure 4: Ablation study. (a) Performance of NePhi trained with various regularizer ratios, indicating that NePhi's performance is insensitive to the regularizer's sampling ratio. (b) and (c) show the registration accuracy and regularity of the transformation field when using NePhi in the optimization-based setting (as in IO, but without predicting the latent code). The registration accuracy and regularity of the transformation field only deteriorates for regularizer ratios below $\sim 4^{-6}$. (d), (e) and (f) show comparisons between using the diffusion regularizer and gradient inverse consistency (NePhi) on NDF (Neural Deformation Field) transformations in the optimization setting.
• Figure 5: The structure of $f_{\theta_3}$ used in the experiments. The number of layers and the dimensions of the feature maps are kept the same for 2D and 3D registration. The position of the $z_l$ branch is adjusted according to the shape of the input image for each registration task to obtain a $z_l$ with a shape close to $16\times 16$ or $16\times 16\times 16$ for 2D and 3D registration respectively.