Dual-stream Pyramid Registration Network¹¹1This work was initially presented at MICCIA 2019.

The goal of 3D medical image registration is to estimate a deformation field $\Phi$ which can warp a moving volume $M\subset R^{H\times W\times D}$ to a fixed volume $F\subset R^{H\times W\times D}$ , so that the warped volume $W=M\circ\Phi\subset R^{H\times W\times D}$ can be accurately aligned to the fixed one $F$. We use $M\circ\Phi$ to denote the application of a deformation field $\Phi$ to the moving volume with a warping operation, with image registration being formulated as an optimization problem:

where $\mathcal{L}_{sim}$ is a function measuring the image similarity between the warped image ($M\circ\Phi$) and the fixed image ($F$), and $\mathcal{L}_{smooth}$ is a regularization constraint on the deformation field ($\Phi$), which enforces spatial smoothness. Both $\mathcal{L}_{sim}$ and $\mathcal{L}_{smooth}$ can be defined in various forms, and recent efforts have been devoted to developing a powerful approach to computing the deformation field $\Phi$. For example, VoxelMorph [4, 5] uses a CNN to compute a deformation field, $\Phi=f_{\theta}(F,M)$, where $\theta$ are learnable parameters of the CNN. In VoxelMorph, the deformation warping operation is implemented by using a spatial transformer network [22], $M\circ\Phi=f_{stn}(M,\Phi)$, and a single-stream encoder-decoder architecture with skip connections (similar to U-Net [38]) is used. Two volumes, $M$ and $F$, are stacked as the input of VoxelMorph. More details of VoxelMorph are described in [4, 5].

Our Dual-PRNet⁺⁺ is built on the encoder-decoder architecture of VoxelMorph, but improves it by introducing a dual-stream design, as shown in Fig. 1. Specifically, the backbone of Dual-PRNet⁺⁺ consists of a two-stream encoder-decoder with shared parameters. We apply an encoder with the same architecture of U-Net [38], which contains five convolutional blocks. Except for the first convolutional block, each block has a 3D down-sampling convolutional layer with a stride of 2, followed by a ReLU operation. Thus the encoder reduces the spatial resolution of the input volumes by a factor of 16 in total, as shown in Fig. 2. In the decoding stage, we apply skip connections to the corresponding convolutional maps in the encoding and decoding process. The lower-resolution convolutional maps (from decoding layers) are up-sampled and concatenated with the higher-resolution ones (from encoding layers), following by a 3$\times$3$\times$3 convolution layer and ReLU operation, as shown in Fig. 1. Finally, we obtain two feature pyramids with multi-resolution convolutional features computed from the moving volume and the fixed volume separately.

The proposed dual-stream design allows us to compute feature pyramids from two input volumes separately, and then predict deformable fields from the learned, stronger and more discriminative convolutional features, which is the key to improve the performance. This is different from existing single-stream networks, such as [4, 5] and [27], which compute the convolutional features from two stacked volumes, and estimate the deformation fields using single-stream convolutional filters. Furthermore, our dual-stream architecture can compute two paired feature pyramids where layer-wise deformation fields can be estimated sequentially at multiple scales. This allows the model to generate a sequence of deformation fields by designing a new sequential pyramid registration method.

Notice that we modify the backbone applied in the original Dual-PRNet [18] by increasing the convolutional blocks from four to five, but reducing the number of channels from 32 at each layer to [8, 16, 16, 32, 32] for the five layers, and also removing the refine units in the original design to keep a lightweight and effective model. This results in a large reduction of the model parameters from 410K to 175K (which might also alleviate the potential overfitting), but maintains the similar performance. For example, it improves the Dice score from 0.631 to 0.653 on the Mindboggle101, but has a reduction with 0.767$\rightarrow$0.743 on the LPBA40. In this work, we will use the new backbone for Dual-PRNet⁺⁺ in all our experiments.

VoxelMorph computes a single deformation field from the convolutional features at the last up-sampling layer in the decoding process, which might make it difficult to handle multi-scale deformations precisely, which are often in case for different anatomical structures of the brain. In this work, we propose a new pyramid registration by designing a set of pyramid registration (PR) modules, which are implemented sequentially at each decoding layer. This allows the model to predict multi-scale deformation fields with increasing resolutions, generating a sequence of pyramid deformation fields, as shown in Fig. 1.

PR module. Each PR module estimates a deformation field at each decoding layer. As input, the PR module uses a pair of convolutional features, together with a deformation field computed from the previous layer (except for the first decoding layer where the deformation field is not available). As output, the PR module yields an estimated deformation field at a given resolution level, which is used in the next pyramid level. The PR module includes a sequence of operations with feature warping, stacking, and convolution (as shown in Fig. 3 (a)), which are implemented repeatedly over the decoding layers.

Sequential operations. Specifically, the first deformation field ($\Phi_{1}$) is computed at the first decoding layer. We first stack the two convolutional features computed at the first decoding layer, and then apply a 3D convolution with size of 3$\times$3$\times$3 to estimate a deformation field. The deformation field ($\Phi_{1}$) is 3D maps with the same shape of the corresponding convolutional feature maps. It is able to extract coarse-level context information, such as high-level anatomical structure of the brain, which is then encoded into the convolutional features computed at the next decoding layer via feature warping: (i) the current deformation field is up-sampled by using bilinear interpolation with a factor of 2, denoted as $u\left(\Phi_{1}\right)$, and (ii) then it is applied to warp the convolutional maps of the moving volume in the next layer, by using a grid sample operation, as shown in Fig. 3 (a). Then the warped convolutional maps are stacked again with the corresponding convolutional features generated from the fixed volume, followed by a convolution operation to estimate a new deformation field. This process is implemented repeatedly at each decoding layer, and can be formulated as,

where $l=1,2,...,N$, indicates the number of decoding layers. $C_{l}^{3\times 3\times 3}$ denotes a 3D convolution at the $l$-th decoding layer. The operator $\circ$ is the warping operation that maps the coordinates of $H^{M}_{l}$ to $H^{F}_{l}$ using $u\left({\ \Phi}_{l-1}\right)$, where $H^{M}_{l}$ and $H^{F}_{l}$ are the convolutional feature pyramids computed from the moving volume and the fixed volume at the $l$-th decoding layer.

Sequential pyramid registration with a set of PR modules was originally introduced in our preliminary version [17]. In this extension, we introduce PR⁺⁺ modules to enhance sequential pyramid registration. It improves the PR module by computing correlation features which are further enhanced by residual convolutions, as shown in Fig. 3 (b). With respect to the PR module, the PR⁺⁺ module includes two additional operations: 3D correlation and residual convolution, which are the key to enhance the learned features and in turn to boost the performance. Specifically, we design a 3D correlation layer to compute correlation features between the warped features (from the moving volume) and the features from the fixed volume (see Fig. 3 (b)). Then the correlation features, together with the two stacked features, are further processed by two convolution blocks with a residual connection to further enhance the representation.

3D correlation layer. In PR⁺⁺ module, a 3D correlation layer is designed to compute the local correlations between the two input volumes in the convolutional feature space. This allows us to aggregate the correlated features which are not directly explored in the original PR module, but can emphasise local details in deep representation.

Specifically, let $P^{W}_{i}$ and $P^{F}_{j}$ denote the central voxel of the 3D blocks (with size of $(2k+1)^{3}$) sampled from the feature maps of the warped moving volume and the fixed volume. The correlation relationship between the two sampled 3D blocks can be computed as:

where $n\in[-k,k]^{3}$ means $n$ iterates over a 3D neighborhood $[-k,k]\times[-k,k]\times[-k,k]$ of $P^{W}_{i}$ or $P^{F}_{j}$. In our experiments, $k$ was empirically set to 1. Given a local 3D block on the feature maps of the (warped) moving volume, it is time-consuming to compute the dense correlations over all the 3D blocks sampled from the feature maps of the fixed volume. Therefore, given a 3D block with $P^{W}_{i}$, we only compute the local correlations by sampling a set of $P^{F}_{j}$ within a 3D neighborhood of $d\times d\times d$, which can be implemented as 3D convolutions. We use a stride $s_{w}=1$ to densely sample $P^{W}_{i}$ from the warped feature maps, and set the correlation neighborhood with $d=3$ on the corresponding fixed feature maps, where $P^{F}_{j}$ is sampled with a stride of $s_{f}=2$. Each sampled block has the same size of $[-k,k]\times[-k,k]\times[-k,k]$ , and we compute direct correlations between two sampled blocks using E.q. (4). This generates 3D correlation maps ($P^{C}$) with shape of $[2\times{FL}(d/s_{f})+1]^{3}\times(H/s_{w})\times(W/s_{w})\times(D/s_{w})$, where $[2\times{FL}(d/s_{f})+1]^{3}=27$ is the number of channels. $FL$ indicates a $Floor$ computation. The generated correlation maps have the same 3D shape as the feature maps of the moving and fixed volumes, which ensure that the three maps can be stacked together for further processing.

Convolutional enhancement. Our dual-stream architecture computes two separate feature pyramids from two input volumes. However, the key to the registration task is to learn the strong anatomical correspondence between the two volumes in the feature space, which inspired us to design a new mechanism to further aggregate the computed pyramid features. The key function of the proposed PR⁺⁺ module is to provide a powerful approach for learning richer local details from the two features, which ensure more accurate estimations of the deformation fields at multiple levels. To enrich the learned features, the computed correlation maps are stacked with the two pyramid features at each decoding layers: the warped features from the moving volume and the pyramid features from the fixed volume, as shown in Fig. 3 (b). The correlation maps have 27 channels over all decoding layers, while the number of channels of the two pyramid features varies over different layers: [8, 16, 16, 32, 32] for the five decoding layers in our experiments.

To this end, we apply two 3D convolution blocks for further processing the stacked features, as shown in Fig. 3 (b). Each convolution block consists of two 3$\times$3$\times$3 convolutional layers, followed by a ReLU operation. The first convolution block reduces the channels of the stacked features considerably from [43, 59, 59, 91, 91] to [8, 16, 16, 32, 32] at the five decoding layers, which are consistent with the numbers of channels applied in the PR module for computational efficiency. In addition, a residual connection is applied to the second convolutional block, in an effort to preserve more context information, and at the same time, to extract discriminative features between the two volumes. Finally, a convolution layer is used to estimate the deformation field. The new PR⁺⁺ module is applied to our dual-stream registration framework, resulting in the enhanced version Dual-PRNet⁺⁺, which boost the performance of the original Dual-PRNet, as demonstrated in our experiments.

The proposed Dual-PRNet⁺⁺ generates five sequential deformation fields with increasing resolutions, indicated as [$\Phi_{1}$, $\Phi_{2}$, $\Phi_{3}$, $\Phi_{4}$, $\Phi_{5}$]. To compute the final deformation field, an estimated deformation field is up-sampled by a factor of 2, and then is warped by the following deformation field being estimated. Such up-sampling and warping operations are implemented repeatedly and sequentially to generate the final deformation field (as shown in Fig. 4), which encodes rich multi-level context information with multi-scale deformations. This allows the model to propagate strong context information over hierarchical decoding layers, where the estimated deformation fields are refined gradually in a coarse-to-fine manner, and thus aggregate both high-level context information and low-level detailed features. The high-level context information equips our model with the ability to work with large-scale deformations, while the fine-scale features allows it to model detailed anatomical structure information. We integrate PR modules or PR⁺⁺ modules into our dual-stream architecture, resulting in an end-to-end trainable model. By simply following VoxelMorph [4, 5], a negative local cross correlation (NLCC) is applied as the loss function, coupled with a smooth regularization, e.g., a diffusion regularizer which computes approximate spatial gradients using differences between neighboring voxels, as detailed in [4].

We compare our Dual-PRNet and Dual-PRNet⁺⁺ with a number of recent approaches: affine registration, SyN [39], VoxelMorph [4, 5], FAIM [27], PMRNet [32], LapIRN [34], Contrastive Registration (CReg) [31] and CycleMorph [24] on both LPBA40 and Mindboggle101 datasets. As discussed in Section 2, VoxelMorph estimates a single deformation field with a single-stream network. FAIM extends VoxelMorph by designing a new penalty loss on negative Jacobian determinants. PMRNet computes average multi-resolution deformation fields, which are obtained from a dual-stream encoder, as the final deformation field. CReg estimates a single transformation field from feature pyramids by using a contrastive loss and a single-stream decoder. Two registration networks were designed in CycleMorph [24], with a cycle consistency, which takes inverse order volumes as inputs. We implemented the affine registration and SyN by using ANTs [3]. For VoxelMorph and FAIM, we used the codes and models provided by the original authors. For CycleMorph, we utilized the released code, and trained the models on the datasets used in our experiment, by using the same experimental settings as [24].

Results and comparisons. On the LPBA40 dataset, as shown in Fig. 7 and Table 1 , Dual-PRNet improves the average Dice score to 0.778, and outperforms SyN, VoxelMorph, and CycleMoprh considerably. With the enhanced PR⁺⁺ module, Dual-PRNet⁺⁺ can further increase the average Dice score by 2.0% (Dice 0.798), and achieves the best performance in the term of average Dice score. However, Dual-PRNet⁺⁺ is outperformed by LapIR in the terms of ASD and HD. With the penalty on both the size and nonsmoothness of the deformation field, LapIR is able to obtain the lowest HD and folding fraction on the determinate of the Jacobian.

The results on the Mindboggle101 are shown in Table 2, where our Dual-PRNet⁺⁺ consistently outperforms the other methods on Dice score, and achieves the best performance on all five regions. It reaches a high average Dice score of 0.748, surpassing the closest one - 0.629 of CReg and CycleMorph, by a large margin. Notice that the new PR⁺⁺ modules lead to a large improvement of $0.631\rightarrow 0.748$ over the original Dual-PRNet, demonstrating its ability to learn detailed brain structure. Furthermore, Dual-PRNet⁺⁺ achieves an ASD of 0.849, which is the best result among all methods, and a comparable HD with LapIRN. This indicates that our method is able to generate less outliers when performing registration. Notice that our methods can achieve excellent Dice scores by simply using CC loss and smooth loss, but did not explored an additional penalty on the deformation fields as did by LapIRN, which would reduce the HD and ASD.

In addition, we also computed folding fractions of Jacobian determinant on deformation fields which measure the smoothness of the deformation fields. Our Dual-PRNet++ obtained a folding faction of 1.725 on the Mindboggle101, outperforming VoxelMorph with 2.274, while FAIM and CycleMorph achieved higher performance of 0.983 and 1.142 respectively. Notice that both FAIM and CycleMorph have a regularity loss designed specifically to encourage the smoothness of the estimated deformation fields in an explicit manner. For example, negative Jacobian determinants were directly measured as a loss in FAIM, while CycleMorph computes a regularization function and a cycle constraint loss to achieve it.

We provide ablation studies to further verify the efficiency of individual technical components developed in our Dual-PRNet⁺⁺. We assessed the benefit of the dual-stream design, sequential pyramid registration with PR modules, and the improved PR⁺⁺ modules. Results from these ablation experiments on the Mindboggle101 are presented in Table 3. To compare our dual-stream architecture with the single-stream design in VoxelMorph, we apply the dual-stream architecture for estimating a single deformation field as VoxelMorph, which achieved an average Dice score of 0.582 on the Mindboggle101 and 0.767 on the LPBA40, improving the single-stream counterpart by +7.1% and +8.4%, respectively. By integrating our sequential pyramid registration with PR modules, the results can be further increased, with +4.9% and +1.1% further improvements on the Mindboggle101 and LPBA40. To further enhance the local details in the learned deep features, we develop the PR⁺⁺ modules which aggregate richer local details by explicitly computing the local correlation features, with residual convolutions for further enhancement. This results in a large further improvement on the Mindboggle101: 0.631$\rightarrow$0.748 on Dice score, as shown in Table 3. Particularly, the residual convolutions and 3D correlation have independent improvements of 0.631$\rightarrow$0.694 and 0.694$\rightarrow$0.748 respectively, making comparable contribution in our design.

We further evaluate the performance of joint segmentation and registration framework by using the proposed Dual-PRNet⁺⁺, which can be integrated into a 3D segmentation network. To make a fair comparison, we replace Voxelmorph with our Dual-PRNet⁺⁺ as the registration network in the joint framework of DeepAtlas [46]. Our Dual-PRNet⁺⁺ estimates a deformation field (the final one as described in Section 2.5) which is then used to warp the available segmentation labels from the source volume to the target one, to guide the learning. Notice that in this implementation, we only optimize the segmentation network by fixing the registration one (also referred as semi-DA in [46]), due to limited GPU memory. Higher performance can be expected with a fully joint learning of the two tasks.

Experiments were conducted on the Mindboggle101. By following DeepAtlas [46], we use 31 labeled regions in the experiments which are different from the 25 regions used in previous registration experiments. Again, with same experimental settings as [46], the joint networks are trained with $N$ labeled volumes and the remained $65-N$ volumes unlabeled. It is a fully supervised learning when $N=65$, which is the total number of the volumes in the dataset. We use $N=21$ in our experiments by following DeepAtlas, and results are compared in Table. 4.

In the case of $N=21$, DeepAtlas with our Dual-PRNet⁺⁺ obtains an average Dice score of 78.04%, clearly outperforming the pure segmentation network (73.48%) which is only trained on 21 labelled volumes. This demonstrates that the joint registration network is greatly helpful to improve the performance of segmentation network, by leveraging the additional unlabelled volumes. Furthermore, as the joint registration network, our Dual-PRNet⁺⁺ can provide more accurate warped anatomical labels than Voxelmorph used by DeepAtlas [46], resulting in large performance improvements on the joint framework, e.g., $61.19\%\rightarrow 66.86\%$ in one-shot learning ($N=1$), and $75.63\%\rightarrow 78.04\%$ when $N=21$. Notice that our result (78.04%) is also closed to the result (81.31%) of fully supervised learning where all labelled volumes ($65$ in total) are used for training. The results suggest that our Dual-PRNet⁺⁺ can work more effectively in the joint segmentation and registration framework, and is helpful to training the segmentation network with limited data and labels provided.

On large displacement. We first visualize the generated multi-resolution deformation fields in Fig. 6 (top). As can be found, a deformation field generated from a lower-resolution layer contains coarse and high-level context information, which is able to encode the high-level semantic information by warping a volume at a larger scale. Conversely, the deformation field estimated from a higher-resolution layer can capture more detailed features. Fig. 6 (bottom) shows the warped images by using the corresponding deformation fields presented. The sequential deformation fields are refined gradually to generate the final field, which can warp the moving image toward the fixed one more accurately, by aggregating more detailed structural information from the preceding fields via sequential warping, as shown in Fig. 6 (c).

We investigate the capability of our methods for handling large spatial displacements, and compare our registration results against that of VoxelMorph in Fig. 8. Our Dual-PRNet and Dual-PRNet⁺⁺ can align the image more accurately than VoxelMorph, especially on the regions containing large spatial displacements, as indicated in green or red regions. In addition, as can be found, Dual-PRNet⁺⁺ has an improvement over the original Dual-PRNet, by using the enhanced PR⁺⁺ modules. The design of 3D correlations with more convolutional layers in the PR++ modules can enlarge the receptive fields, which in turn further enhances the ability to handle large displacements.

To further verify the ability of our Dual-PRNet⁺⁺ to handle large displacements quantitatively, we assume that the large displacements more likely happen on the regions where an affine registration can not perform well, such as the “Occipital” and “Parietal” regions on Mindboggle101, which have low Dice scores of 0.354 and 0.406 respectively. Our Dual-PRNet⁺⁺ can have large relative improvements of 88%-100% in these two regions (where VoxelMorph only has 13%-36% relative improvements), compared to 60%-71% relative improvements on the regions where the affine registration achieves a higher Dice score over 0.450. On the LPBA40, our Dual-PRNet⁺⁺ has relative improvements of 26%-29% on the regions of “Cingulate” and “Temporal” which have low Dice scores of 0.576 and 0.578 by the affine registration, while only achieving about 8% relative improvements on the regions of “Hippocampus” and “Putamen” where the affine registration performs better with 0.753 and 0.775 Dice scores.

On large slice space. We further evaluate the robustness of Dual-PRNet to large slice space. Experiments were conducted on LPBA40, by reducing the slices of the moving volumes from 160$\times$192$\times$160 to 160$\times$24$\times$160. Specifically, we preform the slice reduction on moving volumes by simply removing the slices to 160$\times$24$\times$160, and then interpolate the reduced volumes to the original size (160$\times$192$\times$160) with a spline interpolation (order=1). Then we perform our methods on the reduced-interpolated volumes which have the same size of the original moving volumes. By this way, we can verify different levels of slice reductions between the moving volume and the fixed volume, while keeping the fixed volumes unchanged. During testing, the estimated final deformation field is applied to the labels of the moving volume using zero-order interpolation. With a large reduction of slices from 192 to 24, our Dual-PRNet can still obtain a high average Dice score of 0.711, which even outperforms 0.683 of VoxelMorph [4, 5] using the original non-reduced volumes. This demonstrates the strong robustness of our model against the large spacing displacements.

Compared with LPBA40 dataset, the Mindboggle101 is annotated with the cortical structure, which contains more complicated anatomical structure of the brain, and often requires more accurate local detailed information to identify subtle difference. Fig. 9 demonstrates the registration results on a number of MRI examples from the Mindboggle101. As can be found, VoxelMorph does not provide accurate results on the detailed brain structure. In addition, as demonstrated in ablation studies, our sequential pyramid registration with either PR modules or PR⁺⁺ modules has larger improvements on the Mindboggle101 than that of the LPBA40. Our sequential pyramid registration performs coarse-to-fine refinements of the deformation fields via sequential warping, which naturally aggregate more detailed information from multi-layer feature pyramids. Furthermore, the enhanced Dual-PRNet⁺⁺ can achieve further higher performance by using the improved PR⁺⁺ modules, which are able to compute the local correlations explicitly and thus encodes more detailed structural information. The sequential warping allows the model to propagate the strong high-level context information gradually through the decoding layers, which enhances both the high-level semantic context and the local detailed structure.

We further perform more quantitative analysis. First, we can measure the fineness of the brain structures roughly by computing the ratio of labeled voxels to the total number of voxels in the volumes. The ratios of the labeled voxels on the Mindboggle101 and LPBA40 are 35.83% and 60.02% respectively, which demonstrate that the clinical regions defined in the Mindboggle101 are more fineness. Second, the difficulty of the brain structure on the two datasets can be demonstrated clearly by the performance of an Affine Registration, which has a 0.427 Avg Dice on the Mindboggle101 and a 0.669 Avg Dice on the LPBA40. Therefore, brain structures presented in the Mindboggle101 are more challenging, and our Dual-PRNet++ achieved an over 75% relative improvement, compared to 19% relative improvement obtained on the LPBA40.

We further evaluate the generalization capability of the proposed Dual-PRNet and Dual-PRNet⁺⁺ by conducting external cross-dataset validation. For example, we train on the LPBA40 and test on the Mindboggle101, and vice versa. Dual-PRNet⁺⁺ obtains an average Dice score of 0.788 on the LPBA40, and 0.739 on the Mindboggle101, which are slightly lower than the original performance: 0.798 and 0.748, respectively. Similarly, the cross-data performance of Dual-PRNet are 0.747 and 0.581 on the two databases, compared to the original 0.778 and 0.631 respectively. Therefore, the cross-data performance of both Dual-PRNet and Dual-PRNet⁺⁺ are compared favorably against the original results of VoxelMorph (with 0.683 and 0.511 respectively), demonstrating the improved generalization ability of the proposed methods over different datasets.

Dual-PRNet has its limitation by preforming sequential warping on deformation fields, which might result in an accumulation of interpolation artifacts. Thus it yields the highest folding fraction. However, by integrating our PR⁺⁺ module, the folding fraction of deformation field on Dual-PRNet⁺⁺ can be reduced considerably, reaching a higher performance compared favorably against VoxelMorph. However, we note that the current implementation of Dual-PRNet and Dual-PRNet++ do not have an explicit mechanism to regulate the amount of regularization as it is the case for FAIM [27] or LapIRN [34]. Nonetheless, inspired by these works, we note that adding a regularization term based on the determinant of the Jacobian matrix of the deformation is straightforward in our methods, and worth investigating in line with regularizing the solution where trading accuracy does not affect the downstream clinical task. Besides, we further preformed our Dual-PRNet++ using additional supervision of segmentation masks, which can improve the performance from 0.748 to 0.752 on the Mindboggle101. However, this improvement is relatively limited when compared to that of VoxelMorph. [5]

Besides, in this work, we only performed our methods on 3D brain image (MRI) registration and segmentation. We expect that they can be further applied for or extended to more general 3D medical images, for registration or segmentation, but more experiments and evaluations are required, which can be kept as our future work.