Recurrence With Correlation Network for Medical Image Registration

Voxelmorph is one of the earliest and best known methods for DLIR [3]. It trains a 3D UNet [10] 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 function, $L$, describing the goodness of fit for registering a moving image $m$ to a fixed image $f$ is expressed as:

where $Id$ is the identity transform and $m\circ(Id+D)$ resamples the moving image onto a new grid parametrized by the displacement field, $D$. This resampling operation is carried out differentiably using the spatial transformer network (STN) [11]. The loss function described by (1) 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 [12] and Symmetric Normalization [13].

Laplacian Pyramid Image Registration Network (LapIRN) [5] 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 (1). 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 [5]. Moreover, at coarse resolution the optimization problem is simpler, owing to the fact that the required displacements are smaller in magnitude 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 cost volume layer is used.

The architecture of RWCNet is directly inspired by RAFT, which achieves impressive performance in optical flow [8]. The RAFT network architecture has three key characteristics: 1) CNN feature encoders, 2) cost volume computation at multiple scales between all pairs and 3) recurrent CNN architecture for computing and refining the flow field and performing a ‘lookup’ that subsamples the global cost volume. RWCNet differs from RAFT in several ways. First, computing the cost between all pairs of voxels becomes prohibitively expensive for 3D volumes. Furthermore, computing the cost volume at multiple resolutions simultaneously is also computationally expensive. As such, RWCNet consists of three sub-networks that learn displacement fields at three different resolutions, similar to LapIRN for DLIR or PWCNet[6] for optical flow.

The architecture for the recurrent CNN sub-network is shown in Figure 1. 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, cost volume and moving image are fed into an update block that is a modified gated recurrent unit (GRU) [14]. The GRU is almost identical in implementation to the one used by RAFT [8], 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 cost volume, $C_{1}$, for the next RNN time step. This process of updating displacement field with GRU cell and generating cost volume is repeated for $N$ RNN time steps.

We adopt a course to fine approach to image registration. 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 cost 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.

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