Recurrence With Correlation Network for Medical Image Registration

1. Introduction

Medical image registration describes the task of spatially aligning one medical image to another based on common identifiable features or regions. Medical image registration is commonly required for tasks like image-guided pre-operative planning [1] and quantitative assessments with atlas based segmentation [2], and longitudinal analysis.

Image registration has previously been solved using optimization and variational methods. Deep learning for image registration is an active area of research; deep learning has two potential advantages over existing methods: greatly reduced computational time, and improved accuracy by avoiding local minima. The enhanced speed augments the feasibility of medical image registration for standard clinical applications and offers distinct advantages in time-sensitive settings, including interventional procedures like surgery and focused ultrasound. [3]. Recent investigations have shown improved accuracy of deep learning image registration (DLIR) algorithms over conventional methods [3,4,5].

DLIR, such as Voxelmorph and Laplacian Image Registration Network (LAPIRN), have shown promise in improving the efficiency and accuracy of image registration by leveraging the power of neural networks. However previous methods do not perform well in all scenarios particularly with large displacements and may not adapt to different datasets’ characteristics. Within this study we seek to address this gap by combining three major architectural features: computation of multi-scale features [4], iterative refinement of flow field [6,7] and correlation volume computation [8,9,10] as well as combine these predictions with inference time optimization [11]. To the authors’ knowledge, no prior work has explored a network architecture that combines correlation volume computation, multi-scale refinement and recurrence. This architecture is named recurrence with correlation and abbreviated as RWCNet, for its constituent components. The contributions of this work are:

• A novel network architecture for DLIR, recurrence with correlation network (RWCNet), that outperforms other continuous-domain and multi-scale networks.
• Investigate importance of architectural components (correlation in feature space, multiple scales, and the recurrent structure) in datasets with large displacements.
• Investigate the effect of inference-time optimization had on performance and found varying results.

2. Background

Conventional approaches to medical image registration formulate registration as an energy minimization problem, where the energy functional quantifies the acceptability of any given solution. If t is a spatial transform function that maps a moving image m to a target (or fixed) image, f, a naive approach to finding an optimal registration would be to find a t that maximizes $\mathcal{S}$, a similarity metric that describes the degree of spatial similarity between the moving and fixed images.

$$t^{*}=\underset{t}{max}\mathcal{S}(t(m),f)$$

(1)

where $t^{*}$ is the optimal transform. The moving and fixed images have the same number of dimensions n (i.e., $m,f\in {\mathbb{R}}^n$). In this section, the number of dimensions is assumed to be 3; thus the term for a position in a 3D image, voxel, is used. For simplicity, it can also be assumed that m and f have the same size.

By itself, the above formulation is ill-posed, since for each of the N voxels in the image, there is an unrestricted number of possible sub-voxel locations ($\gg N$) that the transform can map each of the N voxels to. Thus regularization is needed to restrict the search space for the optimal solution. A common framework is variational optimization, where some regularization functional, $\mathcal{R}$, penalizes unlikely transforms based on prior belief. Smooth and small magnitude displacements are examples of prior beliefs that can be represented and penalized explicitly by regularizers. The objective function in this variational optimization scheme becomes:

$$t^{*}=\underset{t}{max}(\mathcal{S}(t(m),f)+\mathcal{R}(t))$$

(2)

Typically, t is a parametric function, belonging to a set of transformations, called the ‘transform model’, T. In the case of rigid registration, where only translation and rotation are permitted operations, T can be the set of functions parametrized by a homogeneous transform matrix with zero scaling and shearing. When the set of transforms is expanded to include affine transformations, the set of allowable functions might be parametrized as transform matrices with non-zero scaling and squaring. Another commonly used transform model is the displacement field, which describes the displacement of each pixel or voxel in the moving image. For 3D medical images

In the case of deformable registration, the transform can be parameterized by a dense grid, $\varphi$. The transformation is then a resampling operation of the moving image onto this new grid. The objective function in this scenario becomes

$${\varphi }^{*}=\underset{\varphi }{max}(\mathcal{S}(m\circ \varphi ,f)+\mathcal{R}(\varphi ))$$

(3)

where $\varphi$ is the displacement field, and $m\circ \varphi$ denotes the resampling operation. A common way to parameterize the displacement field, $\varphi$, is to use a dense deformation field, $D\in {\mathbb{R}}^{3\times H\times W\times D}$, that describes the displacement of each voxel or voxel in a moving image to a fixed image. With this parameterization, the objective function becomes:

$${\varphi }^{*}=\underset{\varphi }{max}(\mathcal{S}(m\circ (Id+D),f)+\mathcal{R}(D))$$

(4)

where $Id$ is the identity grid and $Id+D$ describes the (sub-) voxel coordinates of each moving voxel in the fixed coordinate frame. One strategy to solve such a problem would be through iterative gradient methods. One can perform simple gradient descent on the loss function described above, using:

$${\varphi }^t={\varphi }^{t-1}-\eta {\nabla }_{\varphi }L$$

(5)

where L is the loss function, described by the expression inside the max in Equations (3) and (4), $\eta$ is the learning rate, and ${\nabla }_{\varphi }$ is the jacobian operator.

In terms of inputs and outputs, medical image registration is closely related to optical flow, where the objective of the latter is to compute a flow field describing the motion of objects between two images of, typically, the same scene. Optical flow is an important and well-studied computer vision problem and it has a range of applications including action recognition and pose estimation. Common architectural features of optical flow networks include multi-scale features [12] and cost volume layers [9,12,13].

The computation of multi-scale features and cost volumes for 3D medical image registration have been separately explored in prior work [4,10], which have demonstrated improvements in registration performance on standard datasets. This paper introduces a novel, optical flow-inspired network architecture for DLIR, recurrence with correlation network (RWCNet). RWCNet combines multi-scale iterative features and a correlation volume computation for medical image registration. To the authors’ knowledge, no other work combines these architectural features for DLIR.

3. Related Work

3.1. Voxelmorph

Voxelmorph is one of the earliest and best known methods for DLIR [3]. It trains a 3D U-Net [14] to learn a sub-voxel displacement field, D, that jointly optimizes image fidelity of the transformed image, denoted $\mathcal{S}$, and a regularization loss promoting smooth spatial transformation, denoted $\mathcal{R}$. The loss functional, L, describing the goodness of fit for registering a moving image m to a fixed image f is expressed as:

$$L=\mathcal{S}(m\circ (Id+D),f)+\mathcal{R}(D)$$

(6)

The resampling operation is carried out differentiably using the spatial transformer network (STN) [15]. The loss function described by (6) can be augmented with the correspondence of auxiliary data such as segmentations and keypoints. Voxelmorph was found to be competitive with optimization-based registration methods such as Demons [16] and Symmetric Normalization [17].

3.2. Laplacian Image Registration Network

Laplacian Pyramid Image Registration Network (LapIRN) [4] learns flow fields at $N=3$ resolutions. At each resolution a CNN with residual skip connections learns a displacement field by jointly optimizing an image fidelity term and smoothness term described by (6). LapIRN is trained in a coarse to fine manner; thus the CNN’s at low resolutions are trained before training the networks responsible for learning at higher resolutions. The flow fields learned at low resolutions are upsampled and used to warp the moving input image at higher resolutions.

Multi-scale refinement operates on the principle that an optimal displacement field at low resolution is also a good displacement field at high resolution [4]. Moreover, at coarse resolution the optimization problem is simpler, owing to the fact that the required displacements are a smaller number of the coarse voxels and there are fewer high-level features to match; refining the flow field from a coarse-to-fine resolution can thus simplify the optimization problem. LAPIRN showed improvement in the large-displacement setting, when compared to Voxelmorph and conventional approaches. LapIRN uses convolutional neural network (CNN) architecture with residual connections to learn the flow fields at each resolution.

The method presented in this paper also learns flow fields from coarse to fine resolutions, like LapIRN. RWCNet, the architecture presented in this paper, differs from the one used by LapIRN; a recurrent CNN architecture with features encoders and a correlation volume layer is used.

3.3. ConvexAdam

ConvexAdam [11] combines non-parametric discrete optimization with continuous-domain instance optimization and is able to achieve achieve state-of-the-art performance in several large displacement datasets. RWCNet shares two features with ConvexAdam: they both compute MIND (Modality Invariant Neighborhood Descriptor) features [18] of the fixed and moving image. Second they both use continuous domain optimization as shown in (5) for the final step.

The ConvexAdam algorithm uses an initial discrete optimization whereas RWCNet uses recurrent neural networks to determine the displacement field (discussed below).

3.4. Recurrent Neural Networks for Registration

Recurrent neural networks for image registration have been studied in the 2D [7] and 3D [6] registration settings. In Sandkuhler et al., an RNN composed of a Gated Recurrent Units (GRU) [19] was used to estimate parameters for a parametric Gaussian transformation function. At each time step, the transformation was applied to the moving image, in order to generate a new moving image which was fed into the new time step.

3.5. Recurrent All-Pairs Field Transforms for Optical Flow (RAFT)

Optical flow differs from the general image registration problem because the images are typically of the same object and the consistency of the intensity of objects between both objects is expected. The general image registration problem does not have these constraints. Nevertheless, the two problems are similar. The inputs to the algorithms are two images the output is a flow field from one image to the other. As such, optical flow algorithms can be used directly for medical image registration.

The architecture of RWCNet is inspired by RAFT [13]. The RAFT network architecture has three key characteristics: 1) CNN feature encoders, 2) correlation volume computation at multiple scales between all pairs of pixels in the feature tensor and 3) recurrent CNN architecture for computing and refining the flow field and performing a ‘lookup’ that subsamples the global cost volume. The computation of a correlation volume from feature encoders is a common architectural structure of many optical flow networks [8,9,12] and is based on the notion that the correlation volumes force the network to learn salient features for both input images, because the same object is expected to appear in both images. This might not be happening in other formulations when the inputs are stacked together. The correlation volume (C) can be computed from the feature vectors ($F_f$ and $F_m$) as:

$$C(i,j,k,p,q,r)={\displaystyle \frac{F_f(i,j,k)·F_m(p,q,r)}{\sqrt{n_{\mathrm{features}}}}}$$

(7)

RWCNet differs from RAFT in several ways. RWCNet employs strategies for decreasing the computational complexity. RAFT computes a cost between all pairs of voxels, which is prohibitively expensive for 3D volumes, as such RWCNet instead considers correlation volumes for whole down sampled images, and patches of the image at higher resolutions. Similiarly computing the correlation volume at multiple resolutions simultaneously as is done in RAFT is also computationally expensive. As such, RWCNet considers each resolution in isolation, similiar to the strategy used in LapIRN for DLIR or PWCNet [12] for optical flow.

4. Methods

4.1. Coarse to Fine Registration

We adopt a coarse to fine approach to image registration. Figure 1 provides an overview of this approach. For each resolution, s, a new RNN is trained using inputs from the previous resolution. The weights from previous resolutions are frozen at finer resolutions. At fine resolutions computing the correlation volume for the whole volumes becomes prohibitively expensive; as such, our approach divides the input images into uniform, non-overlapping windows or patches. The size of the patches at each resolution is parameterized by the ‘patch factor’, $p^s\in [0,1]$. The size of the patches at resolution s is computed as $p^s\times S$ where S is the size of the full image at $1\times$ resolution. At higher resolutions, the flow field is used to warp the initial moving image (or patch) using flow fields computed at lower resolutions. Furthermore, the final hidden state is cached at lower resolutions and added to the initial hidden state at higher resolutions, increasing the non-linearity of the network and providing additional context to the network.

4.2. Sub-Network Architecture

The architecture for the recurrent CNN sub-network is shown in Figure 2. Given a fixed and moving image pair (f and m, respectively), the network first learns fixed and moving features ($F_f$ and $F_{m,0}$) by feeding both images through a feature extractor network. A voxel-wise correlation between the fixed and moving features is computed, $C_0$. Due to the large number of dimensions, the correlation is restricted so that only voxels within a certain range, r of the moving voxel are considered. Additionally, the moving image is fed through a context network that extracts contextual information for the hidden network. The output of the contextual network is used as the initial hidden state of the RNN, $h_0$. Finally, a displacement field with zero displacement, $D_0$ is initialized.

The hidden state, flow, correlation volume and moving image are fed into an update block that is a modified gated recurrent unit (GRU) [19]. The GRU is similar to the one implemented in RAFT [13] but in 3D, and it outputs a new hidden state and a new displacement field, $h_1$ and $\Delta D$. The new displacement field is used to update the aggregate displacement, i.e, $D_1=\Delta D+D_0$. This new displacement field is used to warp the moving features (generating $F_{m,1}$), which can be used to generate a new correlation volume, $C_1$, for the next RNN time step. This process of updating displacement field with GRU cell and generating the correlation volume is repeated for N RNN layers.

4.3. Ablation Experiments

Ablation experiments were performed to assess the contribution of the individual architectural features. To test the impact of multi-scale refinement, the network was trained to learn registration fields at a single resolution. To study the impact of correlation, the correlation volume computation was removed; instead the input to the GRU was simply a concatenation of the input features, which was consistent with Voxelmorph and LapIRN.

5. Experiments

5.1. Datasets

Experiments were performed with the OASIS [20] and NLST [21] datasets prepared for the MICCAI 2022 Learn2Reg workshop challenge [22] (Figure 3). The OASIS dataset consists of 414 T1-weighted MRI scans of individuals from ages 18-96 with mild to severe Alzheimer’s. The scans are skull-stripped and resampled onto an isotropic grid and cropped to a uniform size. 35 segmentation labels are provided for important brain regions. The dataset is split into 395 images for training and 19 for validation. Intersubject registration in this context could be used for constructing a sub-population brain atlas or for analysing intensity changes in consistent brain regions that are linked to disease progression.

NLST is a lung-CT dataset with pairs of inhale/exhale scans; keypoints and masks are provided by the Learn2Reg challenge for semi-supervised training. We used a subset of the image pairs (100 out of 150) of the NLST dataset released by the Learn2Reg challenge for training and validation, with a 90:10 training/validation split. Since respiration is accompanied by a large change in lung volume, the displacement field required to register NLST is large, relative to OASIS.

5.2. Training Parameters

The subnetworks at each scale were trained separately from coarse to fine. At higher resolutions, corresponding patches of the fixed and moving images are passed into the network to decrease GPU memory requirements. Table 1 shows the number of steps and the sizes of the inputs at each resolution. To address overfitting, dropout with probability 0.5 was used for the feature networks. The network took about 30 hours to train both on the NLST and OASIS datasets on an NVIDIA A100 with 40GB of RAM.

Distinct composite loss functions were used for the OASIS $L_{\mathrm{OASIS}}$ and NLST $L_{\mathrm{NLST}}$ datasets. An MSE term was used to capture difference between the warped moving image $I_{\mathrm{moving}}$ and the fixed image $I_{\mathrm{fixed}}$ intensities:

$$\mathrm{MSE}={\displaystyle \frac{1}{N}}\sum _{v\in \mathrm{voxels}}{(I_{\mathrm{fixed},v}-I_{\mathrm{moving},v})}^2$$

(8)

The Dice loss term was used to measures the similarity between the warped segmentation $S_{\mathrm{moving}}$ and the fixed segmentation $S_{\mathrm{fixed}}$:

$$L_{\mathrm{DICE}}=1-{\displaystyle \frac{2{\sum }_{v\in \mathrm{voxels}}{S}_{\mathrm{fixed},v}{S}_{\mathrm{moving},v}}{{\sum }_{v\in \mathrm{voxels}}{S}_{\mathrm{fixed},v}+{\sum }_{v\in \mathrm{voxels}}{S}_{\mathrm{moving},v}}}$$

(9)

A regularization loss term used that average magnitude of the gradient of the flow field:

$$\mathcal{R}={\displaystyle \frac{1}{N}}\sum _{v\in \mathrm{voxels}}\left|\nabla {\mathbf{D}}_v\right|$$

(10)

The Total Registration Error (TRE) measures the discrepancy in millimeters between corresponding keypoints in the fixed and moving images. Given m pairs of corresponding keypoints ${\mathbf{p}}_{\mathrm{fixed},i}$ and ${\mathbf{p}}_{\mathrm{moving},i}$ in the fixed and moving images, respectively, the TRE is computed as

$$\mathrm{TRE}={\displaystyle \frac{1}{m}}\sum _{i=1}^m∥{\mathbf{p}}_{\mathrm{fixed},i}-{\mathbf{p}}_{\mathrm{moving},i}∥$$

(11)

For OASIS, MSE was combined with Dice loss and the regularization in a weighted ($w_{\mathrm{MSE}}$, $w_{\mathrm{TRE}}$, and $w_{\mathit{R}}$) sum:

$$L_{\mathrm{OASIS}}=w_{\mathrm{MSE}}·\mathrm{MSE}+w_{\mathrm{DICE}}·L_{\mathrm{DICE}}+w_{\mathcal{R}}·\mathcal{R}$$

(12)

For NLST, the data was range normalized to between 0 and 1, with -4000 serving as the minimum value and 16000 serving as the maximum value. The NLST loss function was a weighted ($w_{\mathrm{MSE}}$, $w_{\mathrm{TRE}}$, and $w_{\mathit{R}}$) summation of the MSE, TRE and regularization term based on the mean gradient of the flow field:

$$L_{\mathrm{NLST}}=w_{\mathrm{MSE}}·\mathrm{MSE}+w_{\mathrm{TRE}}·\mathrm{TRE}+w_{\mathcal{R}}·\mathcal{R}$$

(13)

For both datasets, an Adam optimizer with learning rate of $3\times {10}^{-4}$ was used.

LAPIRN, another multi-resolution model, was trained on OASIS using training parameters from [4]. We use the non-diffeomorphic variant, which does not ensure topological consistency of the spatial transformation, but achieves greater quantitative accuracy. For NLST, we augment training with a supervised discrepancy loss between the fixed and moving keypoints once the displacement is applied to the moving keypoints. We use a MSE loss instead of the normalized cross correlation (NCC) loss prescribed by the original paper to be consistent with the the loss used for training RWCNet.

6. Results

Figure 4 shows qualitative results generated for the NLST and OASIS datasets. Table 2 summarizes the results when comparing RWCNet with LAPIRN on the NLST and OASIS datasets. RWCNet outperforms LAPIRN on both datasets. However, the difference in Dice is quite modest 0.7%. For the NLST dataset, the difference in performance is much more pronounced, with the difference in average TRE being >3mm. These results suggest that the architectural features of RWCNet, significantly aid generalization performance in the more challenging large-displacement setting.

These ablation results (Table 3) provide interesting insights into the role that architecture plays in registration accuracy in different datasets. Unsurprisingly, multi-resolution registration plays a crucial role in the accuracy of RWCNet; registering at only $4\times$ downsampling on the NLST dataset yields a keypoint discrepancy of 5.52 mm, whereas registering at multiple resolutions yields a discrepancy of 2.11mm. OASIS Dice, likewise, drops from 80.7% to 74.0%.

The impact of correlation and number of RNN time steps was markedly different for both datasets. In the OASIS datasets, replacing correlation with stacking of the input feature tensors did not drastically impact the registration performance. The Dice score dropped by 0.7. Likewise, when only 2 RNN time steps were used in RWCNet, the drop in accuracy was even lower in the OASIS dataset; the Dice score only dropped by 0.6%. This was in contrast to NLST, where decreasing the number of RNN time steps and removing correlation drastically decreased performance. The keypoint accuracy decreased to 4.10mm when correlation was not computed. Likewise, when only 2 time steps were used, the accuracy decreased to 5.17mm. Instance optimization, which fine tuned the flow field using an ADAM optimizer, did improve the performance of RWCNet in NLST dataset so that it was on par with ConvexAdam (that also uses instance optimization). Interestingly the performance of RWCNet was degraded in the OASIS dataset by instance optimization, again to a performance that was on par with ConvexAdam.

The intrinsic dataset attributes(inter vs intra subject or large vs small displacements), might have contributed to the observed variations in performance within the ablation study. OASIS has relatively small displacements compared to NLST and as such may not benefit as much from correlation volumes and the RNN structure. In contrast both datasets showed that a multi-resolution approach improved performance. These observations could be helpful for designing generalizable methods for image registration problems.

7. Conclusions

The DLIR network developed in this investigation, RWCNet, borrows key ideas from optical flow literature; specifically that of correlation volume computation and iterative refinement, to improve medical image registration performance. For the large displacement NLST dataset, results with instance optimization were the same with either RWCNet or ConvexAdam implying that the steps other than instance optimization did not influence the final output for those setups. Without instance optimization it was found that RWCNet was more performant than comparable deep learning methods. In the small displacement setting, RWCNet was found to have a very small improvement over LAPIRN another deep learning approach and improved on improved over variational approaches. More tuning of algorithm parameters is need if instance optimization is to be incorporated into registration algorithms with deep learning. Future work should investigate whether the instance optimization step can be augmented with features learned by the network to ensure a more uniform improvement in all datasets.

Ablations illustrated the benefits of the adopted methods. Multi-scale iterative refinement played a key role; performing registration at a single resolution drastically decreases the accuracy of the model. RNN steps and the computation of the explicit correlation volume were notably influential, their primary effect was observed in the dataset with large displacements. On OASIS, the Dice score did not change drastically without these features. This suggests that these architectural features were valuable in large displacement settings. This was a significant finding, as it can inform future work for large displacement registration.

The findings in this paper highlight the need for dataset-specific hyperparameter selection. In the small displacement setting RWCNet yields only a modest improvement over LapIRN. In the larger displacement setting, RWCNet drastically outperformed LapIRN. RWCNet was more resource intensive than LapIRN therefore there was little benefit to using RWCNet for cases with small displacements. The findings of this work were consistent with the concept of automatic self-configuring networks that vary their architecture components based on the dataset. Such an approach for image registration would be similar to nnU-net [23] for medical image segmentation, this would be in contrast to variational image registration that have been distributed as one size fits all with parameters adjustable by the user. The approach has many potential application where fast computation is needed such as intra-procedural registration, further there are applications where variational approaches do not converge to the global optimal registration trasform.