Unsupervised 3D End-to-End Medical Image Registration with Volume Tweening Network

TL;DR

This paper introduces Volume Tweening Network, an unsupervised end-to-end CNN framework for 3D medical image registration that is significantly faster and achieves state-of-the-art accuracy without manual annotations.

Contribution

The paper presents a novel unsupervised CNN-based registration method with an end-to-end cascading scheme, integrated affine registration, and invertibility loss for improved accuracy and efficiency.

Findings

880x faster than traditional methods

Achieves state-of-the-art registration accuracy

Operates efficiently without manual annotations

Abstract

3D medical image registration is of great clinical importance. However, supervised learning methods require a large amount of accurately annotated corresponding control points (or morphing), which are very difficult to obtain. Unsupervised learning methods ease the burden of manual annotation by exploiting unlabeled data without supervision. In this paper, we propose a new unsupervised learning method using convolutional neural networks under an end-to-end framework, Volume Tweening Network (VTN), for 3D medical image registration. We propose three innovative technical components: (1) An end-to-end cascading scheme that resolves large displacement; (2) An efficient integration of affine registration network; and (3) An additional invertibility loss that encourages backward consistency. Experiments demonstrate that our algorithm is 880x faster (or 3.3x faster without GPU acceleration) than traditional optimization-based methods and achieves state-of-theart performance in medical image registration.

Tables

Table 1. Table I: Ratio of loss functions.

Subnetwork	Loss	Relative Ratio
Affine	Similarity	1
	Determinant	0.1
	Orthogonality	0.1
Dense 1	Similarity	0
	Total variation	1
Dense 2	Similarity	0.05
	Total variation	1
Dense 3	Similarity	1
	Total variation	1

Table 2. Table II: Comparison among traditional methods, VoxelMorph and our VTN (liver).

Method	Seg. IoU	Lm. Dist.	Time	Time (w/o GPU)
ANTs [4]	0.8124	11.93	N/A	748 s
Elastix [5]	0.8365	12.36	N/A	115 s
VoxelMorph-2 [23]	0.6796	18.10	0.20 s	17 s
VTN ADDD	0.8868	12.04	0.13 s	26 s
VTN ADDD + inv	0.8882	11.42	0.13 s	26 s
supervised	0.7680	13.38	0.13 s	26 s

Table 3. Table III: Standard deviation of Jacobian determinant and fraction of folding on the liver dataset. Areas with negative Jacobian determinant are considered folding. Standard deviations across instances are in parentheses.

Method	Std. Jacobian	Folding (%)
ANTs [4]	0.241 (0.030)	0.194 (0.145)
Elastix [5]	0.264 (0.072)	0.001 (0.006)
VTN ADDD	0.787 (0.171)	0.365 (0.268)
VTN ADDD + inv	0.770 (0.169)	0.288 (0.240)

Table 4. Table IV: Comparison of performance of VTN ADDD with different amount of unlabeled training data (liver).

Training Dataset for VTN ADDD	Seg. IoU	Lm. Dist.
LITS	0.8712	13.11
BFH	0.8815	12.78
LITS + BFH	0.8868	12.04

Table 5. Table V: Comparison of performance with different number of cascaded subnetworks (liver).

Network	Seg. IoU	Lm. Dist.	Time	Time (w/o GPU)
D	0.8119	14.44	0.08 s	10 s
DD	0.8556	12.97	0.10 s	20 s
DDD	0.8709	12.49	0.12 s	28 s
AD	0.8323	13.20	0.09 s	9 s
ADD	0.8703	12.28	0.11 s	19 s
ADDD	0.8868	12.04	0.13 s	26 s

Table 6. Table VI: Comparison among traditional methods, VoxelMorph [ 23 ] and our algorithm on brain datasets. Bold: best among all methods. Star: best among methods with ANTs affine pre-alignment. Methods with “ ⋆ ” use ANTs for affine pre-alignment. The preprocessing time (about 74 seconds) is not included in the table.

Method	Seg. IoU	Lm. Dist.	Time	Time (w/o GPU)
ANTs [4]	0.9387	2.85	N/A	764 s
Elastix [5]	0.9180	3.23	N/A	121 s
VoxelMorph-2⋆ [23]	0.9268	2.84	0.19 s	14 s
VTN DD⋆	0.9305⋆	2.83⋆	0.09 s	19 s
VTN ADD	0.9270	2.62	0.10 s	17 s
VTN ADD + inv	0.9278	2.64	0.10 s	17 s
supervised	0.9060	2.94	0.10 s	19 s

Table 7. Table VII: Standard deviation of Jacobian determinant and fraction of folding on the brain dataset. Areas with negative Jacobian determinant are considered folding. Standard deviations across instances are in parentheses.

Method	Std. Jacobian	Folding (%)
ANTs [4]	0.230 (0.037)	0.294 (0.228)
Elastix [5]	0.171 (0.034)	0.000 (0.000)
VTN ADD	0.339 (0.126)	0.018 (0.043)
VTN ADD + inv	0.332 (0.122)	0.017 (0.055)

Table 8. Table VIII: Comparison of performance of VTN DD (with ANTs affine alignment) with different amount of unlabeled training data (brain).

Training Dataset for VTN DD	Seg. IoU	Lm. Dist.
ADNI + ABIDE-1	0.9299	2.90
ADNI + ABIDE-1 + ABIDE-2	0.9312	2.86
ADNI + ABIDE-1 + ABIDE-2 + ADHD	0.9305	2.83

Table 9. Table IX: Comparison of performance with different number of cascaded subnetworks (brain), and comparison between using ANTs as affine alignment and end-to-end network (integrated affine registration subnetwork). The first group (with “ ⋆ ”) uses ANTs to affinely align input images, whereas the second group does not. Bold: best among all methods. Star: best among methods with ANTs affine pre-alignment. Methods with “ ⋆ ” use ANTs for affine pre-alignment. The preprocessing time (about 74 seconds) is not included in the table.

Method	Seg. IoU	Lm. Dist.	Time	Time (w/o GPU)
D⋆	0.9241	2.91	0.07 s	10 s
DD⋆	0.9305	2.83⋆	0.09 s	19 s
DDD⋆	0.9320⋆	2.85	0.10 s	28 s
AD	0.9214	2.73	0.08 s	9 s
ADD	0.9270	2.62	0.10 s	17 s
ADDD	0.9286	2.63	0.11 s	26 s

Equations

$$I_1\left(x\right)\approx I_2\left(x+f\left(x\right)\right),$$

$$\mathrm{warp}\left(\mathrm{warp}\left(I,g_1\right),g_2\right)\left(x\right)$$

$$=$$

$$g_1\star g_2=g_2+\mathrm{warp}\left(g_1,g_2\right),$$

$$\mathrm{warp}\left(\mathrm{warp}\left(I,g_1\right),g_2\right)=\mathrm{warp}\left(I,g_1\star g_2\right).$$

$$\mathrm{Cov}\left[I_1,I_2\right]=\frac{1}{\left|\Omega \right|}\sum _{x\in \Omega }{I}_1\left(x\right)I_2\left(x\right)-\frac{1}{{\left|\Omega \right|}^2}\sum _{x\in \Omega }{I}_1\left(x\right)\sum _{y\in \Omega }{I}_2\left(y\right),$$

$$\mathrm{CorrCoef}\left[I_1,I_2\right]=\frac{\mathrm{Cov}\left[I_1,I_2\right]}{\sqrt{\mathrm{Cov}\left[I_1,I_1\right]\mathrm{Cov}\left[I_2,I_2\right]}}.$$

$$L_{\mathrm{Corr}}\left(I_1,I_2\right)=1-\mathrm{CorrCoef}\left[I_1,I_2\right].$$

$$L_{\mathrm{TV}}=\frac{1}{3|\Omega |}\sum _x\sum _{i=1}^3(f(x+e_i)-f(x){)}^2,$$

$$L_{\mathrm{ortho}}=-6+\sum _{i=1}^3\left({\lambda }_i^2+{\lambda }_i^{-2}\right).$$

$$L_{\det }={\left(-1+\det \left(A+I\right)\right)}^2,$$

$$f\left(x\right)=Ax+b.$$

$$L_{\mathrm{inv}}={\Vert f_{12}\star f_{21}\Vert }_2^2+{\Vert f_{21}\star f_{12}\Vert }_2^2.$$

Full text

I Introduction

Image registration is the process of finding the non-linear spatial correspondence between the input images (see Figure 1). It has a wide range of applications in medical image processing, such as aligning images of one subject taken at different times. Another example is to match an image of one subject to some predefined coordinate system, such as an anatomical atlas [1].

Over recent decades, many traditional algorithms have been developed and studied to register medical images [2, 3, 4, 5]. Most of them define a hypothesis space of possible transforms, often described by a set of parameters, and a metric of the quality of such a transform, and then find the optimum by iteratively updating the parameters. Traditional methods have achieved good performance on several datasets, and are state-of-the-art, but their registration speed is barely practical for clinical applications. These methods do not exploit the patterns that exist in the registration task. In contrast, learning-based methods are generally faster. Computationally, they do not need to iterate over and over before producing the result. Conceptually, they learn those patterns represented by the parameters of the underlying learning model. The trained model is an efficient replacement of the optimization procedure.

However, there has not been an effective approach to generate ground-truth fields for medical images. Widely used supervised methods require accurately labeled ground truth with a vast number of instances. The quality of these labels directly affects the result of supervision, which entails much effort in traditional tasks such as classification and segmentation. But flow fields are dense and ambiguous quantities that are almost impossible to be labeled manually, and moreover, automatically generated dataset (e.g., the Flying Chairs dataset [6]) which deviates from the realistic demands is not appropriate. Consequently, supervised methods are hardly applicable. In contrast, unlabeled medical images are universally available, and sufficient to advance the state of the art through our unsupervised framework shown in this paper.

Optical flow estimation is a closely related problem that aims to identify the correspondence of pixels in two images of the same scene taken from different perspectives. FlowNet [6] and its successor FlowNet 2.0 [7] are CNNs that predict optical flow from input images using fully convolutional networks (FCN [8]), which are capable of regressing pixel-level correspondence. FlowNet is trained on Flying Chairs, a synthetic dataset that consists of images generated using computer graphics algorithms and the ground-truth flows.

Spatial Transformer Networks (STN) [9] is a component in neural networks that spatially transforms feature maps to ease back-end tasks. It learns a localization net to produce an appropriate transformation to “straighten” the input image. The localization net is learnt without supervision, though back-end task might be supervised. Given a sampling grid and a transformation, STN applies the warping operation and outputs the warped image for further consumption by deeper networks. The warping operation deals with off-grid points by multi-linear interpolation, hence is differentiable and can back-propagate gradients.

Inspired by FlowNet and STN, we propose Volume Tweening Network (VTN), which enables the unsupervised training of end-to-end CNNs that perform voxel-level 3D medical image registration (see Figure 2). The moving image is registered and warped, and the warped image is compared against the fixed image to form a similarity loss. There is a rich body of research into similarity losses [10, 11, 12]. The model is trained to minimize a combination of regularization loss and similarity loss. As the method is unsupervised, the performance potentially increases as it is trained with more unlabeled data. The network consists of several cascaded subnetworks, the number of which might vary, and each subnetwork is responsible for producing a transform that aligns the fixed image and the moving one. Deeper layers register moving images warped according to the output of previous layers with the initial fixed image. The final prediction is the composition of all intermediate flows. It turns out that network cascading significantly improves the performance in the presence of large displacement between the input images. While the idea of cascading subnetworks is found in FlowNet 2.0, our approach does not include as much artificial intervention of network structures as FlowNet 2.0 does. Instead, we employ a natural dichotomy in a subnetwork structure consisting of affine and deformable registration processes, which is also found in traditional methods, including ANTs (affine and SyN for deformable) [4] and Elastix (affine and B-spline for deformable) [5]. Besides these structural innovations, we also introduce the invertibility loss to 3D medical image registration, which encourages backward consistency while achieving better accuracy. Compared with traditional optimization-based algorithms, ours is 880x faster (or 3.3x faster without GPU acceleration) and achieves state-of-the-art performance.

To summarize, we present a new unsupervised end-to-end learning system using convolutional neural networks for deformation field prediction between 3D medical images. In this framework, we develop 3 technical components: (1) We cascade the registration subnetworks, which improves performance for registering largely displaced images without much slow-down; (2) We integrate affine registration into our network, which proves to be effective and faster than using a separate tool; (3) We incorporate an additional invertibility loss into the training process, which improves registration performance. The contributions of this work are closely related. An unsupervised approach is very suitable for this problem, as the images are abundant and the ground truth is costly to acquire. The use of the warping operation is crucial to our work, providing the backbone of unsupervision, network cascading and invertibility loss. Network cascading further allows us to plug in different subnetworks, and in this case, the affine registration subnetwork and the deformable ones, enabling us to adapt the natural structure of multiple stages in image registration. The efficient implementation of our algorithm gives a satisfactory speed. The proposed VTN is also used as a building block in [13], which thoroughly exploits the idea of deep cascades and the code is publicly available111https://github.com/microsoft/Recursive-Cascaded-Networks.

II Related Work

II-A Traditional Algorithms

Previously, there has been much effort on automating image registration. Tools like FAIR [14], ANTs [4] and Elastix [5] have been developed for automated image registration [1]. Generally, these algorithms define a space of transformations and a metric of alignment quality, and then find the optimal transformation by iteratively updating the parameters. The optimization process is highly time-consuming, rendering such methods impractical for clinical applications.

The transformation space can be either parametric or non-parametric. Affine transforms can be described by only a few real numbers, whereas a free-form dense deformable field specifies the displacement for each grid point. Though the latter contains all possible transforms, it is common to apply a multi-stage approach to the problem. For example, ANTs [4] registers the input images with a rigid transform, then a general affine transform, and finally a deformable transform modeled by SyN. In this paper, we also have components for affine/deformable registration, which are modeled by neural networks.

Selecting an informative metric is crucial to registration quality. The metrics, also called loss functions, often consist of two parts, one measuring the level of correspondence between the input images implied by the transform, the other regularizing the transform itself. Examples of the former part include photo-metric difference, correlation coefficient and mutual information [11, 12] among others [10]. Some of these measures, notably mutual information, require binning or quantizing, which makes gradients vanish thus continuous optimization techniques such as gradient descent inapplicable.

II-B Supervised Learning Methods

Lee et al. [15] employ Support Vector Machines (SVM) to learn the similarity measure for multi-modal image registration for brain scans. The training process of their algorithm requires pre-aligned image pairs thus the method is supervised. There have been some works tackling medical image registration with CNNs [16]. Sokooti et al. [17] develop a patch-based CNN to register chest CT scans and trains it with synthetic data. Miao et al. [18] use CNN to perform 2D/3D registration, in which CNN regressors are used to estimate transformation parameters. Their CNNs are trained over synthetic data and the method is not end-to-end. FlowNet [6], developed by Dosovitskiy et al., is an FCN [8] for optical flow prediction. The network is trained with synthetic data and estimates pixel-level correspondence.

While supervised learning methods achieve good performance, either abundant groud-truth alignment must be available, or synthetic data are used. Generation of synthetic data has to be carefully designed so that the generated data resemble the real ones.

II-C Unsupervised Learning Methods

In an earlier work towards fast and accurate medical image registration by Shan et al. [19], an unsupervised end-to-end learning-based method for deformable medical image registration is proposed. The method in [19] registers 2D images. It is evaluated by registering corresponding sections of MR brain images and CT liver scans. Some later work directly performs 3D registration [20, 21], fully exploiting the information available from the 3D images. They aim to learn sparse parameters introduced by the traditional algorithms, which probably limits their performance. de Vos et al. [21] also try to stack multiple networks and investigate an affine network, however, each of them is trained independently and separately. Their method is not end-to-end and only achieves a comparable performance to their baseline methods. Another work [22] proposes to train a multi-step affine network followed by a momentum generator network. Although their framework is claimed to be “end-to-end”, each stage is still trained separately after fixing the previous ones. In contrast, with the help of jointly trained cascaded subnetworks, the affine stage is naturally integrated into our framework, rather than done or trained out-of-band in [19, 20, 21, 22].

VoxelMorph, proposed by Balakrishnan et al. [23, 24], is an unsupervised learning-based method for 3D medical image registration that predicts a dense deformation field. Another recently proposed unsupervised method by Krebs et al. [25, 26], based on a low-dimensional probablistic model, performs comparably to VoxelMorph. VoxelMorph contains an encoder-decoder structure, uses warping operation to produce warped moving images, and is trained to minimize the dissimilarity between the warped image and the fixed image. Their method does not consider affine registration and assumes the input images are already affinely aligned, whereas ours embeds the process of affine registration as an integrated part of the network. Furthermore, their algorithm does not work well when large displacement between the images is present, which is common for liver CT scans. Finally, VoxelMorph is designed to take any two images as input, but [23] only evaluates it in the atlas-based registration scenario. Consider a clinical scenario where the brain of a patient is captured before and after an operation. It would be better, in terms of accuracy and convenience, to register these two images, instead of registering both to an atlas.

It is noticeable that all the unsupervised methods use the warp operation to train networks without supervision. This paper exploits the operation in a trio (namely enabling unsupervised training, enabling cascading, and implementing invertibility loss, which will be detailed later). We present a more sophisticated and robust design that works well in the presence or absence of large displacement, and we evaluate the methods for general registration among images.

III Method

III-A Problem Formulation

The input of an image registration problem consists of two images $I_{1,2}$, both of which are functions $\Omega\to\mathbb{R}^{c}$, where $\Omega$ is a region of $\mathbb{R}^{n}$ and $c$ denotes the number of channels. Since this work focuses on 3D medical image registration, we confine ourselves to the case where $\Omega$ is a cuboid and $c=1$ (grayscale image). Specifically, this means $\Omega\subseteq\mathbb{R}^{3}$ and each image is a function $\Omega\to\mathbb{R}$. The objective of image registration is to find a displacement field (or flow field) $f_{12}:\Omega\to\mathbb{R}^{3}$ so that

[TABLE]

where the precise meaning of “$\approx$” depends on specific application. The field $f_{12}$ is called the flow from $I_{1}$ to $I_{2}$ since it tells where each voxel in $I_{1}$ is in $I_{2}$. We define $\mathrm{warp}\left({I_{2},f}\right)$ as the image $I_{2}$ warped according to $f$, i.e., $\mathrm{warp}\left({I_{2},f}\right)\left(x\right)=I_{2}\left({x+f\left(x\right)}\right)$. The above objective can be rephrased as finding $f$ maximizing the similarity between $I_{1}$ and $\mathrm{warp}\left({I_{2},f}\right)$.

The image $I_{1}$ is also called the fixed image, and $I_{2}$ the moving one. The term “moving” suggests that the image is transformed during the registration process.

Consider warping an image twice, first with $g_{1}$ then with $g_{2}$. What this procedure produces is

[TABLE]

This motivates the definition of the composition of two flows. If we define the composition of the flow fields $g_{1},g_{2}$ to be

[TABLE]

Equation (III-A) can be restated as

[TABLE]

It is noticeable that the warp operation in the above formulation should be further specified in practice. Real images as well as flow fields are only defined on lattice points. We continuate them onto the enclosing cuboid by trilinear interpolation as done in [27]. Furthermore, we deal with out-of-bound indices by nearest-point interpolation. That is, to evaluate a function defined on lattice points at any point $x$, we first move $x$ to the nearest in the enclosing cuboid of those lattice points, then interpolate the value from the 8 nearest lattice points.

III-B Unsupervised End-to-End Registration Network

Our network, called Volume Tweening Network (VTN), consists of several cascaded registration subnetworks, after each of which the moving image is warped. The unsupervised training of network parameters is guided by the dissimilarity between the fixed image and each of the warped images, with the regularization losses on the flows predicted by the subnetworks.

The warping operation, also known as the sampler in STN [27], is differentiable to both the input image and the input flow field due to the trilinear interpolation. Those warping operations can back-propagate gradients to all the preceding subnetworks, which is critical for the end-to-end learning of our cascaded networks.

In deformable image registration, it is common to apply an initial rigid transformation as a global alignment before predicting the dense flow field. Instead of prepending a time-consuming preprocessing stage with a tool like ANTs [4] as done in VoxelMorph [23], we integrate this procedure as a top-level subnetwork. Our affine registration subnetwork predicts a set of affine parameters, after which a flow field can be generated for warping. The integrated affine registration subnetwork not only works in negligible running time, but also outperforms the traditional affine stage.

III-C Loss Functions

To train our model in an unsupervised manner, we measure the (dis)similarity between the moving images warped by the spatial transformer and the fixed image. Regularization losses are introduced to prevent the flow fields from being unrealistic or overfitting. We use correlation coefficient as the similarity measurement, and the total variation loss as the regularization term for dense flow predictions. Furthermore, we introduce orthogonality loss and determinant loss as the regularization terms for the affine registration subnetwork, which are essential in preventing the gradients from exploding and ensuring training stability. Those loss functions are discussed as follows.

Correlation Coefficient

The covariance between $I_{1}$ and $I_{2}$ is defined as

[TABLE]

where $\Omega$ denotes the cuboid (or grid) on which the input images are defined. Their correlation coefficient is defined as

[TABLE]

The images are regarded as random variables whose sample space is the points on which voxel values are available. The range of correlation coefficient is $\left[{-1,1}\right]$, it measures how much the two images are linear related, and attains $\pm 1$ if and only if the two are linear function of each other. Applying a non-degenerate linear function to any of the images does not change their correlation coefficient, therefore, this measure is more robust than $L_{2}$ loss. For real-world images, the correlation coefficient should be non-negative (unless one of the images is a negative film). The correlation coefficient loss is defined as

[TABLE]

Total Variation Loss (Smooth Term)

For a dense flow field, we regularize it with the following loss that discourages discontinuity:

[TABLE]

where $e_{1,2,3}$ form the natural basis of $\mathbb{R}^{3}$. This varies from the initial definition of the total variation loss [28] (which includes a square root), but our formula is more natural as a loss term (referring to the L2 regularization).

Orthogonality Loss

For the specific task discussed in this paper (medical image registration), it is usually the case that the input images need only a small scaling and a rotation before they are affinely aligned. We would like to penalize the network for producing overly non-rigid transform. To this end, we introduce a loss on the non-orthogonality of $I+A$, where $I$ denotes the identity matrix and $A$ denotes the transform matrix produced by the affine registration network (see Section IV-B for more details). Let $\lambda_{1,2,3}$ be the singular values of $I+A$, the orthogonality loss is

[TABLE]

The motivation of this formula is mainly that a matrix is orthogonal if and only if all its singular values are $1$. Hence, the more deviant $I+A$ is from being an orthogonal matrix (i.e., the more its singular values deviating from ones), the larger its orthogonality loss. If $I+A$ is orthogonal, the value will be zero.

Computing orthogonality loss involves singular values of $I+A$. The square of those singular values are exactly the eigenvalues of ${\left({I+A}\right)}^{\mathrm{T}}{\left({I+A}\right)}$. Since the loss is a symmetric function of those eigenvalues, it can be rewritten as a fraction w.r.t. the coefficients of the characteristic polynomial of ${\left({I+A}\right)}^{\mathrm{T}}{\left({I+A}\right)}$ by Viète’s theorem. Then the derivatives could be directly calculated.

Determinant Loss

We assume images are taken with the same chirality, therefore, an affine transform involving reflection is not allowed. This imposes the requirement that $\det\left({I+A}\right)>0$. Together with the orthogonality requirement, we set the determinant loss to be

[TABLE]

where $I$ denotes the identity matrix and $A$ denotes the transform matrix produced by the affine registration network.

IV Network Architecture

IV-A Cascading

Each subnetwork is responsible for aligning the fixed image and the current moving image. Following each subnetwork, the moving image is warped with the predicted flow, and the warped image is fed into the next cascaded subnetwork. The flow fields are composed to produce the final estimation. Figure 2 illustrates how the networks are cascaded, how the images are transformed and how each part contributes to the loss. With all layers being differentiable, gradients can back-propagate to the whole system and enable the unsupervised end-to-end learning.

It might be tempting to compare our scheme with that of FlowNet 2.0 [7]. FlowNet 2.0 stacks subnetworks in a different way than our method. It performs two separate lines of flow estimations (large/small displacement flows) and fuses them into the final estimation. Each of its intermediate subnetworks has inputs of not only the warped moving image and the fixed image, but also the initial moving image, the current flow, and the brightness error. Its subnetworks are carefully crafted, having similar yet different structures and expected to solve specific problems (e.g., large/small displacement) in flow estimation. In contrast, our method does not involve two separate lines of registration, i.e., each subnetwork works on the fixed image and the warped moving image produced by the previous one, and intermediate subnetworks do not get more input than the initial one. We do not interfere much with the structures of subnetworks. Despite the initial affine subnetwork, we do not assign specific tasks to the remaining subnetworks since they share the same structure.

IV-B Affine and Dense Deformable Subnetworks

The affine registration subnetwork aims to align the input image with an affine transform. It is only used as our first subnetwork. As illustrated in Figure 3, the input is downsampled by strided 3D convolutions, and finally a fully-connected layer is applied to produce $12$ numeric parameters as output, which represents a $3\times 3$ transform matrix $A$ and a $3$-dimensional displacement vector $b$. As a common practice, the number of channels doubles as the length of resolution halves. The flow field produced by this subnetwork is defined as

[TABLE]

The moving image is then transformed according to the output affine parameters and fed into the subsequent deformable registration subnetworks.

The dense deformable registration subnetwork is used as all subsequent subnetworks, each of which refines the registration based on the output of the subnetwork preceeding it. It follows an encoder-decoder architecture, as illustrated in Figure 4, which is commonly used for dense prediction. We use strided 3D convolution to progressively downsample the image, and then use deconvolution (transposed convolution) [8] to recover spatial resolution. As suggested in U-Net [29], skip connections between the convolutional layers and the deconvolutional layers are added to help refining dense prediction. The subnetwork will output the dense flow field, a volume feature map with $3$ channels ($x,y,z$ displacements) of the same size as the input.

IV-C Invertibility

Given two images $I_{1,2}$, going from a voxel in $I_{1}$ to the corresponding voxel in $I_{2}$ then back to $I_{1}$ should give zero displacement. Otherwise stated, the registration should be round-trip. In Figure 5, we demonstrate the possible situations. The pair of solid arrows exemplifies round-trip registration, whereas the pair of dashed arrows exemplifies non-round-trip registration. If we have computed two flow fields (back and forth), $f_{12}$ and $f_{21}$, the composed fields exhibit the round-trip behavior of the registration, as illustrated by the magenta straight arrow in Figure 5. Ideally, round-trip registration should satisfy the equations $f_{12}\star f_{21}=f_{21}\star f_{12}=0$. We capture the round-tripness for a pair of images with the invertibility loss, namely

[TABLE]

The larger the invertibility loss, the less round-trip the registration. For perfectly round-trip registration, the invertibility loss is zero. We come up with, formulate, and implement the invertibility loss independently of [30]. We use L2 invertibility loss whereas [30] uses L1 left-right disparity consistency loss, which is just a matter of choice. We are the first to incorporate the invertibility loss into 3D images to boost performance on medical image tasks.

V Experiment

We evaluate our algorithm with extensive experiments on both liver CT datasets and brain MRI datasets. The models are trained separately for the two types of images. We train and test the models for pairwise, subject-to-subject registration. This setting is more general than the atlas-based registration as studied in VoxelMorph [23], where all moving images are assumed to be registered to the fixed atlas.

The input images to the models are assumed to be of size $128^{3}$. We apply spatial augmentation during training, where the input images are transformed by random B-Spline [31] fields of $5\times 5\times 5$ control points within a maximum displacement of $12$.

We compare our algorithm against state-of-the-art traditional registration algorithms including ANTs222The command we use for ANTs: -d 3 -o <OutFileSpec> -u 1 -w [0.025,0.975] -r [<Fixed>,<Moving>,1] -t Rigid[0.1] -m MI[<Fixed>,<Moving>,1,32,Regular,0.2] -c [2000x2000x2000,1e-9,15] -s 2x1x0 -f 4x2x1 -t Affine[0.1] -m MI[<Fixed>,<Moving>,1,32,Regular,0.1] -c [2000x2000x2000,1e-9,15] -s 2x1x0 -f 4x2x1 -t SyN[0.15,3.0,0.0] -m CC[<Fixed>,<Moving>,1,4] -c [100x100x100x50,1e-9,15] -s 3x2x1x0 -f 6x4x2x1 [4] and Elastix333The command we use for Elastix: -f <Fixed> -m <Moving> -out <OutFileSpec> -p Affine -p BSpline_1000 [5], as well as VoxelMorph [23]. Our algorithm achieves state-of-the-art performance while being much faster. Our experiments prove that the performance of our unsupervised method is improved as more unlabeled data are used in training. We show that cascading subnetworks significantly improves the performance, and that integrating affine registration into the method is effective.

We evaluate the performance of algorithms with the following metrics:

• •

Seg. IoU is the Jaccard coefficient between the warped liver segmentation and the ground truth. We warp the segmentation of the moving image by the predicted deformable field, and compute the Jaccard coefficient of the warped segmentation with the ground-truth segmentation of the fixed image. “IoU” means “intersection over union”, i.e., $\frac{\left|{A\cap B}\right|}{\left|{A\cup B}\right|}$, where $A,B$ are the set of voxels the organ consists of.

• •

Lm. Dist. is the average distance between warped landmarks (points of anatomical interest) and the ground truth.

• •

Time is the average time taken for each pair of images to be registered. Some methods are implemented with GPU acceleration, therefore there are two versions of this metric (with or without GPU acceleration).

Our model is defined and trained using TensorFlow [32]. We accelerate training with nVIDIA TITAN Xp and CUDA 8.0. We use the Adam optimizer [33] with the default parameters in TensorFlow 1.4. The batch size is 8 pairs per batch. The initial learning rate is ${10}^{-4}$ and halves every epoch after the $4^{\mathrm{th}}$ epoch. Each epoch consists of 20000 batches and the number of epochs is 5. Performance evaluation uses the same GPU model. Traditional methods work on CPU. We also test neural-network-based methods with GPU acceleration disabled for a fairer comparison of speed. The CPU model used is Intel® Xeon® CPU E5-2690 v4 @ 2.60GHz (14 Cores).

V-A Experiments on Liver Datasets

V-A1 Settings

The input to the algorithms are liver CT scans of size ${128}^{3}$. Affine subnetworks downsample the images to $4^{3}$ before applying the fully-connected layer. Dense deformable subnetworks downsample the images to $2^{3}$ before doing transposed deconvolution.

We cascade up to 4 registration subnetworks. The reason we need to cascade multiple subnetworks is that large displacement is very common among CT liver scans and that with more subnetworks, images with large displacement can be progressively aligned. Among these networks, the one with one affine registration subnetwork and three dense deformable registration networks (referred to as “ADDD”) is used to be trained with different amount of data and compared with other algorithms. We use the correlation coefficient as our similarity loss, the orthogonality loss and the determinant loss as regularization losses for the affine subnetwork, and the total variation loss as that for dense deformable subnetworks. The ratio of losses for “ADDD” is listed in Table I. The performance is not very sensitive to the choice of hyper-parameters, since each of the dense deformable subnetworks can automatically learn to progressively align the images, and only the final subnetwork and the affine subnetwork need to be trained with similarity loss.

V-A2 Datasets

We have three datasets available:

• •

LITS [34] consists of 130 volumes. LITS comes with segmentation of liver, but we do not use such information. This dataset is used for training.

• •

BFH is provided by Beijing Friendship Hospital and consists of 92 volumes. This dataset is used for training.

• •

MICCAI (MICCAI’07) [35] consists of 20 volumes with liver segmentation ground truth. We choose 4 points of anatomical interest as the landmarks444The landmarks: (L1) the top point of hepatic portal; (L2) the intersection of the superior and anteroir branches of the right lobe; (L3) the intersection of the superior and inferior branches of the right lobe; and (L4) the intersection of the medial and inferior branches of the left lobe. and ask 3 expert doctors to annotate them, taking their average as the ground truth. This dataset is used as the test data.

All pairs of images in the mixture of LITS and BFH are sampled into mini-batches for training. All image pairs in MICCAI are tested during evaluation. We crop raw liver CT scans to a volume of size ${128}^{3}$ around the liver, and normalize them by adjusting exposure so that the histograms of the images match each other. The preprocessing stage is necessary as the images come from different sources.

V-A3 Comparison among Methods

In Table II, “ADDD” is our model detailed in Section V-A1, and “ADDD + inv” is that model trained with additional term of invertibility loss in the central area (the beginning and the ending quaters of each side are removed) with relative weight ${10}^{-3}$. Learning-based methods (VTN and VoxelMorph) are trained on LITS and BFH datasets. All methods are evaluated on the MICCAI dataset. To prove the effectiveness of our unsupervised method, we also train “ADDD” supervised (the row “supervised”), where the output of ANTs is used as the ground truth (using end-point error [6] plus regularization term as the loss function).

VoxelMorph-2 is trained (using code released by [23]) with a batch size of 4 and an iteration count of 5000. Its performance with a batch size of 8 or 16 is worse so a batch size of 4 is used. After 5000 iterations, the model begins to overfit. The reason that it does not perform well on liver datasets might be that it is designed for brain registration thus cannot handle large displacement.

The results in Table II show the vast speed-up of learning-based methods against optimization-based methods. The Wilcoxon signed-rank test indicates that our methods significantly surpass state-of-the-art registration algorithms in terms of both Segmentation IoU and Landmark Distance.

Table III presents the smoothness metrics (as done in [21]) of our method compared to the traditional methods. It is not surprising that our significantly better performance comes at the price of worse deformation smoothness, but the fraction of folding is still well acceptable (less than 1%). The invertibility loss has a considerable impact on reducing the fraction of folding.

Figure 6 compares different methods listed in Table II, where three landmarks are selected and the sections of the volumes at the height of each landmark in the fixed image are rendered. This means the red crosses (landmarks in the moving and warped images) indicate the projections of the landmarks onto those planes. It should be noted that though the sections of the warped segmentations seem to be less overlapping with those of the fixed one, the Segmentation IoU is computed for the volume and not the sections. It might well be the case that the overlap is not so satisfactory when viewed from those planes yet is better when viewed as a volume. Similarly, overlapping red and yellow crosses do not necessarily imply overlapping fixed and warped landmarks as they might deviate along $z$-axis.

In Figure 7, we compare ADDD + inv, ADDD, ADD, AD and D. It shows that network cascading better aligns the images with the presence of large displacement, while the invertibility loss has a remarkable effect on the liver boundary.

V-A4 Performance with Different Amount of Data

In Table IV, the “ADDD” network (see Section V-A1) trained with full data (LITS + BFS) is compared with that trained with only a part of data. The result demonstrates that training with more unlabeled data improves the performance. Since we do not need any annotations during the training phase, there are abundant clinical data ready for use.

V-A5 Network Cascading

In Table V, all networks are trained on LITS + BFH. The models whose name does not include “A” have the affine registration subnetwork removed, and the number of “D”s is the number of dense deformable registration subnetworks. (See Section V-A1 for “ADDD”.)

V-B Experiments on Brain Datasets

Settings

The input to the algorithms are brain MR images of size ${128}^{3}$. For some experiments, the selection of which will be detailed later, the brain scans are preprocessed to be aligned to a specific atlas from LONI [36]. We use ANTs for this purpose. There are two reasons we align the brain scans with ANTs. One is that VoxelMorph [23] requires the input to be affinely registered. The other is that we will compare the performance between “ANTs + deformable registration subnetworks” and “affine registration subnetwork + deformable registration subnetworks”, i.e., to compare the effectiveness of integrating our affine registration subnetwork in place of ANTs.

The following methods will have ANTs-affine-aligned images as input: VoxelMorph methods and VTN without “A” (i.e., “D”, “DD” and “DDD”). A more precise naming is “ANTs + VoxelMorph” and “ANTs + VTN”. The following methods will not have ANTs-affine-aligned images as input: ANTs, Elastix, VTN with “A” (i.e., “AD”, “ADD” and “ADDD”). Those methods have affine registration integrated. The comparison inside the former group focuses on dense deformable registration performance, that inside the latter group on overall performance, and that among the two groups benchmarks affine registration subnetwork versus ANTs affine registration in the context of a complete registration pipeline.

We will show 3 sets of comparisons similar to those for liver datasets. In the tables listed in later paragraphs, the time (with or without GPU) does not include the preprocessing with ANTs even if it is used. Preprocessing an image with ANTs costs 73.94 seconds on average. We will mention this fact again when such emphasis is needed.

Care should be taken when evaluating methods with ANTs affine alignment. For the data to be comparable with those with affine registration integrated, the fixed image should be equivalent. Methods with ANTs affine alignment have both moving and fixed images aligned to an atlas. Those with integrated affine registration never move the fixed image. The affine transform produced by ANTs might not be orthogonal, which is the source of unfair comparison. If the affine transform is shrinking, methods with ANTs affine alignment gain advantage. Otherwise, methods with integrated affine alignment do.

One measure, Segmentation IoU, is not affected, because the volumes of all objects get multiplied by the determinant of the affine transform and the evaluation measure is homogeneous. For Landmark Distance, we perform the inverse of the linear part of the affine transform (which aligns the fixed image to the atlas) to the difference vector between warped landmark and landmark in the (aligned) fixed image, so that the length goes back to the coordinate defined by the original fixed image. This way, we minimize loss of precision to prevent unfairly underevaluating methods with ANTs affine alignment. Speaking of the actual data, the affine transformations produced by ANTs are slightly shrinking. Our correction restores a fair comparison among all methods.

Datasets

We use volumes from the following datasets for training:

• •

ADNI [37] (67 volumes);

• •

ABIDE-1 [38] (318 volumes): part of data from ABIDE;

• •

ABIDE-2 [38] (1101 volumes): the rest from ABIDE;

• •

ADHD [39] (973 volumes).

We acquire the second part of ABIDE after a while when the first part was downloaded and processed, thus the split. This only helps us to understand how performance improves as more data are used for training. For comparison among different methods, it is always the case that all the data mentioned above are used for training. All pairs of images in the mixture of those datasets are sampled into mini-batches for training.

Raw MR scans are cropped to ${128}^{3}$ around the brain. The scans are normalized based on the histograms of their foreground color distribution, which might vary because they are captured on different sites.

For evaluation, we use 20 volumes from the LONI Probabilistic Brain Atlas (LPBA40) [36]. LONI consists of 40 volumes, 20 of which have tilted head positions and are discarded. For the remaining 20 volumes, 18 landmarks555The landmarks: (L1) right lateral ventricle posterior; (L2) left lateral ventricle posterior; (L3) anterior commissure corresponding to the midpoint of decussation of the anterior commissure on the coronal AC plane; (L4) right lateral ventricle superior; (L5) right lateral ventricle inferior; (L6) left lateral ventricle superior; (L7) left lateral ventricle inferior; (L8) middle of lateral ventricle; (L9) posterior commissure corresponding to the midpoint of decussation; (L10) right lateral ventricle superior; (L11) left lateral ventricle superior; (L12) middle of lateral ventricle; (L13) corpus callosum inferior; (L14) corpus callosum superior; (L15) corpus callosum anterior; (L16) corpus callosum posterior tip of genu corresponding to the location of the most posterior point of corpus callosum posterior tip of genu on the midsagittal planes; (L17) corpus callosum fornix junction; and (L18) pineal body. are annotated by 3 experts and the average are taken as the ground truth.

Comparison Among Methods

In Table VI, we compare different methods on brain datasets. All neural networks are trained on all available training data, i.e., ADNI, ABIDE and ADHD. In the table, “supervised” is our “ADD” model supervised with ANTs as the ground truth. Its loss is the end-point error [6]. Table VII presents the smoothness metrics of our method compared to the traditional methods.

Among these methods, our “ADD” achieves the lowest Landmark Distance with a competitive speed. If we compare “ADD” with “DD”, the Wilcoxon signed-rank test indicates that the integration of affine registration subnetwork significantly improves the Landmark Distance, compared to using ANTs for out-of-band affine alignment. Although there is still a performance gap compared to ANTs in terms of Segmentation IoU, it is remarkable that our method generates significantly less folding area.

Figure 8 exemplifies the methods on three pairs of scans. Comparison between our methods and traditional methods proves the applicability of our methods to 3D brain registration. Comparison between ADD and DD⋆ shows that integrating affine registration subnetwork is effective.

Performance with Different Amount of Data

In Table VIII, we summarize the performance of “DD” (with ANTs affine alignment) trained on different amount of unlabeled data. As more data are used to train the network, its performance in terms of Landmark Distance consistently increases.

Network Cascading and Integration of Affine Registration

Table IX compares the performances of differently cascaded networks. The networks without “A” have ANTs-affine-aligned images as input, whereas the networks with “A” do not.

As one would expect, the performance in each group improves as the model gets more levels of cascaded subnetworks. While the methods with ANTs affine alignment have higher Segmentation IoU, integrating affine registration subnetwork yields better Landmark Distance. Worth mentioning is that the better Segmentation IoU comes at the price of a rather slow preprocessing phase (74 seconds).

VI Discussion

Our method achieves significant performance gain over traditional methods on the liver CT dataset, however, some method (ANTs) still performs slightly better on the brain MRI dataset. This is probably because brain MRIs are mainly of small displacements which might be better suited for traditional iterative methods. Besides the correlation coefficient loss used in this paper, other similarity measurements like cross correlation can also be explored for better accuracy.

VII Conclusion

In this paper, we present Volume Tweening Network (VTN), a new unsupervised end-to-end learning framework using convolutional neural networks for 3D medical image registration. The network is trained in an unsupervised manner without any ground-truth deformation. Experiments demonstrate that our method achieves state-of-the-art performance, and that it witnesses an 880x (or 3.3x without GPU acceleration) speed-up compared to traditional medical image registration methods. Our thorough experiments prove our contributions, each on its own being useful and forming a strong union when put together. Our networks can be cascaded. Cascading deformable subnetworks tackles the difficulty of registering images in the presence of large displacement. Network cascading also enables the integration of affine registration into the algorithm, resulting in a truly end-to-end method. The integration proves to be more effective than using out-of-band affine alignment. We also incorporate the invertibility loss into the training process, which further enhances the performance. Our methods can potentially be applied to various other medical image registration tasks.

Acknowledgment

The authors would like to thank Beijing Friendship Hospital for providing the “BFH” dataset as well as other data providers for making their data publicly available.