Medical Image Registration via Neural Fields

TL;DR

A new neural network based image registration framework, called NIR (Neural Image Registration), which is based on optimization but utilizes deep neural networks to model deformations between image pairs and achieves better results in shorter computation times.

Abstract

Image registration is an essential step in many medical image analysis tasks. Traditional methods for image registration are primarily optimization-driven, finding the optimal deformations that maximize the similarity between two images. Recent learning-based methods, trained to directly predict transformations between two images, run much faster, but suffer from performance deficiencies due to model generalization and the inefficiency in handling individual image specific deformations. Here we present a new neural net based image registration framework, called NIR (Neural Image Registration), which is based on optimization but utilizes deep neural nets to model deformations between image pairs. NIR represents the transformation between two images with a continuous function implemented via neural fields, receiving a 3D coordinate as input and outputting the corresponding deformation vector. NIR provides two ways of generating deformation field: directly output a displacement vector field for general deformable registration, or output a velocity vector field and integrate the velocity field to derive the deformation field for diffeomorphic image registration. The optimal registration is discovered by updating the parameters of the neural field via stochastic gradient descent. We describe several design choices that facilitate model optimization, including coordinate encoding, sinusoidal activation, coordinate sampling, and intensity sampling. Experiments on two 3D MR brain scan datasets demonstrate that NIR yields state-of-the-art performance in terms of both registration accuracy and regularity, while running significantly faster than traditional optimization-based methods.

Figures (8)

• Figure 1: Overview of NIR, which is a optimization-based pairwise medical image registration framework via neural fields. In each iteration of optimization, every position $p$ is sampled from the coordinate of target volume $T$ and the deformed position $p^{\prime}$ is predicted by NF. The intensity similarity loss $\mathcal{L}_{sim}$ between sampled image intensities is govern by the local normalized cross-correlation and the regularization term $\mathcal{L}_{jedt}$ penalizes the regions where the local deformation orientations are inconsistent, as formulated in Eq. \ref{['eq:loss']}. During inference, the neural field takes as input the whole grid and outputs the deformed positions of the whole grid. Then, by sampling the intensity of the deformed grid on the moving volumes, we can get the warped volumes. Plot (b) only presents the transformation of moving volumes via NIR, but the structures associated with the moving volumes can also be transformed in the same way.
• Figure 2: Neural Fields for Coordinate Deformations -- In the above figure, blue modules indicate the parameters to be optimized. (a) illustrates the neural deformation field that directly transforms the coordinate $p$ in the target volume to the coordinate $p^{\prime}$ in the moving volume. (b) illustrate the neural velocity field which predicts the stationary velocity vector along the deformation trajectory from $p$ to $p^{\prime}$. The neural velocity field plays as the dynamic function of a NODE solver and the final deformations are obtained via the integration of the predicted velocity vector.
• Figure 3: Coordinate Samplers and Performance Comparisons. (a), (b) and (c) illustrate sampling 16 coordinates per batch from total 64 2D coordinates with three kinds of coordinate samplers. (d) ranks the registration performance of NIR models optimized with two practical coordinate samplers (downsize sampler and mini-patch sampler) in four aspects. The higher ranking in each dimension indicates better performance in that aspect. As is shown in (d), consuming almost the same GPU memory during optimization, compared to NIR optimized with the mini-patch sampler, NIR optimized with the downsize sampler can take less time to converge to a more accurate registration results with more violations in topology preserving. The expected solution, as indicated by the red-dot line, should be of great performance in both registration accuracy and regularity with no or modestly extra computations. For the numerical results supporting the ranking in plot (d), please refer to Tab. \ref{['tab:comp_among_cs']}.
• Figure 4: Overview of NIR with Hybrid Coordinate Sampling Scheme. The optimization is composed of two phases, in which two neural fields ($NF_1$ and $NF_2$) are optimized separately. In the first phase, $NF_1$ is optimized with the downsize coordinate sampler ($CS_1$) for 200 iterations. In the second phase, with the mini-patch coordinate sampler ($CS_2$), fixed $NF_1$ provides the initial deformations and only $NF_2$ is optimized. During inference, NIR with hybrid coordinate sampler requires grid coordinates to pass through two neural fields in sequence to get the deformed coordinates.
• Figure 5: Comparison of different coordinate samplers (a), (b) and (c) are registration results of NIR optimized with the downsize sampler, mini-patch sampler and hybrid NIR. The above image pair are 'OASIS_OAS1_0001_MR1' ($T$) and 'OASIS_OAS1_0002_MR1' ($M$) from the OASIS dataset and we present the registration results over the optimization iterations, generated by the differomorphic NIR. DSC and $J_{\leq 0}$ are the evaluation metrics for registration accuracy and regularity separately. Details about the dataset and evaluation metrics can be found in Sec. \ref{['sec:exp']}.