Multi-scale Neural ODEs for 3D Medical Image Registration

The optimization problem for image registration in Eq. 1 can be solved via gradient descent optimizer. Thus, the update of $\theta$ at each step is $\theta_{t+1}=\theta_{t}-\eta_{t}{\partial\mathcal{(L+R)}}/{\partial\theta_{t}}\triangleq\theta_{t}+f(\theta_{t},t)$, where $t$ represents $t^{th}$ step and $\eta_{t}$ is the step size. If the time difference is small enough, the difference equation can be re-written as an ODE function, $\frac{d\theta_{t}}{dt}=f(\theta_{t},t),t\in[0,T]$. Given the initial parameter $\theta_{0}$, the final parameter $\theta_{T}$ is the solution to this ODE initial value problem. Some conventional methods such as LDDMM [6] and SyN [3] describe the evolution of transformation as a differential equation. These methods solve differential equations where system dynamics are described by predefined functions, which is less flexible. Neural ODEs use trainable network to replace the predefined function, which can be considered as learning an optimizer to compute the gradient update. $\theta_{T}$ is the inference of the neural ODE model. We further assume that for two 3D images $x_{a}$ and $x_{b}$, the evolution of $\theta_{t}$ follows a neural ODE:

where $f_{w}$ is a neural network with parameter $w$ which takes the current warped image $x_{a}\circ\phi_{\theta_{t}}$, the target image $x_{b}$ and the time variable $t$ as inputs. Therefore, the final output at time $T$ can be computed by integrating function $f$ over time interval $[0,T]$. In practice, the integral is evaluated by an ODE solver, such as Runge–Kutta methods and adaptive step size solver. To train the neural ODE model, we adopt the adaptive checkpoint adjoint method [33].

Multi-scale ODE network: Empirically, we found that solving neural ODE for image registration problem requires an extensive number of function evaluations (NFE) by ODE solver, which leads to prolonged training and inference time. To address this problem, we propose a multi-scale ODE network (MS-ODENet) which performs registration by solving ODEs at different resolutions. Specifically, let $\{x_{i}^{l}\}_{l=1}^{L}(i=a,b)$ be an image pyramid with $L$ different resolutions, where $x^{(L)}_{i}=x_{i}$, and $x^{(l-1)}_{i}$ is generated by down sampling $x^{(l)}_{i}$ with a factor of 2 on all 3 axes. We divide the whole time interval $[0,T]$ into $L$ congruent segments. In each segment, we solve the ODE in Eq. 3 at the corresponding resolution as shown in Fig. 2 a and b.

where $f^{(l)}_{w_{l}}$ is the network at the $l$-th scale. The output parameter $\theta_{T}$ is the integral over all scales. The benefits of this design are two-fold: 1) The time cost for function evaluations is much smaller at low resolutions, allowing a larger number of steps to reach the desired accuracy. 2) The searching space for image registration is largely reduced and therefore the convergence of the ODE network is faster, which also makes the model less sensitive to local optimal.

Loss functions: The proposed registration network is trained in an unsupervised manner by minimizing a loss function similar to Eq. 1. The utilized similarity metric $\mathcal{L}_{\text{sim}}$ (Fig. 1 b) between two images ($x_{a}$ and $x_{b}$) is defined in Eq. 4:

where $E^{c}_{{}_{3D}}(\cdot)$ is a 3D content feature extractor. The 3D content feature is composed of 2D content features generated from $N$ randomly selected 2D images of different axes from the image utilizing a 2D feature extractor $E^{c}$ (Section 2.2). In addition, we perform a random perturbation to a given image $x_{a}$ such that $\tilde{x}_{a}=x_{a}\circ\phi_{\tilde{\theta}}$, where $\tilde{\theta}$ is randomly sampled from a distribution of parameter $\mathbb{P}_{\theta}$. The pair $x_{a}$ and $\tilde{x}_{a}$ are fed to the registration network resulting transformation parameters $\tilde{\theta}_{T}=\text{MS-ODENet}(\theta_{0},x_{a},\tilde{x}_{a},L,T)$. Our network can align two images from the same modality so that we expect $\tilde{\theta}_{T}$ approximating to the purturbation $\tilde{\theta}$ by minimizing L2 loss: $\mathcal{L}_{\text{self}}=\mathbb{E}_{x_{a},a,\tilde{\theta}}||\tilde{\theta}-\tilde{\theta}_{T}||_{2}^{2}$. The total loss function is summarized as $\mathcal{L}=\mathcal{L}_{\text{sim}}+\lambda_{\text{self}}\mathcal{L}_{\text{self}}+\lambda_{\text{reg}}\mathcal{L}_{\text{reg}}$, where $\mathcal{L}_{\text{reg}}=\mathbb{E}_{x_{a},x_{b},a,b}||\nabla\phi_{\theta_{T}}||_{2}^{2}$ is the regularization term enforcing a smooth motion field in deformable registration. $\lambda_{\text{self}}$ and $\lambda_{\text{reg}}$ are weighting coefficients.

Different contrast images from the same patient share the same anatomical structures (content features) but have different style features. This inspires us decomposing the image into content features and style features in the latent space. The realization of this feature domain disentanglement is extended from the diverse image-to-image translation framework [19]. We trained content encoder $E^{c}$, style encoder $E^{s}$ and generator $G$ to perform image translation from one contrast to another. Note that only $E^{c}$ is used in our registration task. Fig. 2c illustrates an overview of our feature extraction framework utilizing image-to-image translation. Suppose two different groups $\mathcal{X}_{a}$ and $\mathcal{X}_{b}$ sampled from $M$ modality groups ($\{\mathcal{X}_{i}\}_{i=1}^{M}$). Given $x_{a}$ and $x_{b}$ two image samples in these two groups, we perform cross-modality translation as follows: (a) Extract the content and style features, i.e., $s_{a}=E^{s}(x_{a},a)$, $c_{a}=E^{c}(x_{a})$, $s_{b}=E^{s}(x_{b},b)$, $c_{b}=E^{c}(x_{b})$. (b) Swap style features and generate translated images, i.e., $\hat{x}_{ab}=G(c_{a},s_{b},b)$, $\hat{x}_{ba}=G(c_{b},s_{a},a)$. (c) Encode the translated images to reconstruct content and style features of the original images, $\hat{s}_{a}=E^{s}(\hat{x}_{ba},a)$, $\hat{c}_{a}=E^{c}(\hat{x}_{ab})$, $\hat{s}_{b}=E^{s}(\hat{x}_{ab},b)$, $\hat{c}_{b}=E^{c}(\hat{x}_{ba})$. (d) Swap the style features again to generate the original images, $\hat{x}_{aba}=G(\hat{c}_{a},\hat{s}_{a},a)$, and $\hat{x}_{bab}=G(\hat{c}_{b},\hat{s}_{b},b)$. We trained the network using adversial loss in  [19]. Moreover, we expect that $\hat{x}_{aba}$ and $\hat{x}_{bab}$ are consistent to the original images $x_{a}$ and $x_{b}$ respectively. We use L1 loss $\mathcal{L}_{\text{cyc}}$ to enforce this cycle consistency. We also introduce the feature reconstruction loss for content/style features, $\mathcal{L}_{\text{rec}}^{c}$/$\mathcal{L}_{\text{rec}}^{s}$, which are the L1 loss between $c_{a},c_{b}$/$s_{a},s_{b}$ and $\hat{c}_{a},\hat{c}_{b}$/$\hat{s}_{a},\hat{s}_{b}$. Similarly, we define mono-modal reconstruction loss $\mathcal{L}_{\text{rec}}^{x}$ as the L1 loss between the reconstructed images $\hat{x}_{a}=G(c_{a},s_{a},a)$, $\hat{x}_{b}=G(c_{b},s_{b},b)$ and the original images $x_{a},x_{b}$. Since training is not trivial in 3D due to high dimensionality, we use a 2D multi-channel (adjacent slices in 3D volume) network.

Dataset: Currently, there is a lack of large scale medical image dataset dedicated to multi-modal image registration. We use a public dataset and a separate acquired volunteer dataset. The public dataset is the brain tumor segmentation (BraTS) 2020 dataset [22] which consists of multi-modal 3D brain MRI with four distinctive contrasts (T1, T2, T2-FLAIR, and T1Gd) from 494 subjects with glioblastomas, resulting in various multi-modal image registration tasks (12 pairs per subject). Besides, tumor masks provide a clinically meaningful evaluation metric for registration. For each subject, all modalities were normalized into 3D volumes with a size of $240\times 240\times 160$ and a resolution of $1\times 1\times 1\text{mm}^{3}$. The dataset is split into 444 subjects (5,328 pairs) for training and validation and 50 subjects (600 pairs) for testing. Since different modalities have been registered, we simulate motions by applying random rigid transformation, random control point deformation, or both. Note that the simulated transformation fields are only used for evaluation, not for training. We also acquired additional multi-modal 3D MR brain data on 25 volunteers using a special MR technique that can acquire multiple contrasts simultaneously. This private dataset is used only for testing the generalizability of the proposed methods. The motion is simulated as described above. Images were preprocessed to remove skull using DeepBrain¹¹1https://github.com/iitzco/deepbrain and resized as the public data.

Registration networks: We evaluated our methods on rigid, deformable and a hybrid of rigid and deformable motions. For 3D rigid motion, $f$ includes convolution layers followed by fully connected layers. The output size of $f$ is the same as the degree of freedom in the transformation. For deformable motion, fully convolutional network is used. One variation of the deformable registration is kernel based method (B-spline). Parameter $\theta$ is the grid of control points which can also be regarded as an image. For dense motion, the UNet as in  [26] is used for voxel-wise estimation. For hybrid motion, we cascade the rigid and deformable MS-ODENets sequentially.

Feature extraction network: The backbone network is similar to [15]. The content encoder is a fully convolutional network while the style encoder has convolution layers followed by a global average pooling and a fully connected layer, which forces the network to extracts global style features. The generator is a CNN with deconvolution layers to generate images of the original size. Besides, to condition on modality $\mathcal{X}_{i}$ for style encoder and image generator, the modality code $i$ is converted into a one-hot vector concatenated with the input tensor along the channel dimension.

Implementation details: We follow the protocol in [2] for training GAN with number of slices set to 5. For registration, we set $\lambda_{\text{self}}=1$, and let $\lambda_{\text{reg}}$ be 10 and 2 for dense and B-spline respectively. We use $L=T=3$ for MS-ODENet. Let $F(dt)$ be Euler’s method with fixed step size $dt$ and $A(\epsilon)$ be adaptive Heun’s method with tolerance of error $\epsilon$. We use adaptive solvers when the search space is small (Rigid: $A(10^{-3})$-$A(10^{-3})$-$F(0.1)$) to avoid extensive NFE and use fixed step size solver for large search space (B-spline: $F(0.2)$-$F(0.2)$-$F(0.25)$, Dense: $F(0.2)$-$F(0.2)$-$F(0.5)$). All networks are trained using Adam [17] optimizer with a learning rate of $10^{-4}$ on an NVIDIA Tesla V100 GPU.

Evaluation metrics: For evaluation, Dice scores [12] are computed over tumor masks. With available synthetic transformation fields and ground truth images, root mean square errors are calculated and denoted as RMSE($\phi$) and RMSE($x$), respectively. The metrics are averaged over all pairs of test data.

Table 1 shows the quantitative results for rigid, deformable and hybrid registration (rigid+deformable). Rigid: Random rotation and translation were synthesized along all three axes and were sampled from $U(-75^{\circ},75^{\circ})$ and $U(-20,20)\mbox{mm}$, respectively. We used the rigid registration method with MI metric in Advanced Normalization Tools (ANTs) [4] as the baseline. We also used the trained GAN to perform image translation followed by mono-modal image registration with ANTs (ANTs+I2I). Deformable: We made synthetic image pairs through elastic transformations by perturbing B-spline control points with noise from $\mathcal{N}(0,(8\text{mm})^{2})$ on three axes. For comparison, we use SyN [3] in ANTs, VoxelMorph (VM) [5] with MI metric, and also their variants with image translation (ANTs+I2I and VM+I2I). For MS-ODENet, we use two different parameterizations, namely dense deformation (D) and B-spline (B). Fig. 3 shows example results. Rigid + deformable: Random rotations, translations and control point noise are from $U(-40^{\circ},40^{\circ})$, $U(-10,10)\mbox{mm}$, and $\mathcal{N}(0,(8\text{mm})^{2})$, respectively. We compared our method to SyN, RCN [32] and VM. RCN is an iterative deep learning method consisting of an affine network and $n$ deforamble networks. We use MI metric for training RCN and set $n=2$ due to memory limit. Since VM was proposed to solve deformable registration, we performed a rigid registration using our rigid MS-ODENet prior to applying VM. This approach is denoted as VM^∗.

Ablation: Fig. 4(a) summarizes the results that evaluate models with no learned content loss (replaced by MI), no multi-scale ODE, no self-supervision or no multi-slice GAN respectively and compare them with the full model on rigid + B-spline deformable registration. To investigate the necessity of iterative registration for large motions, we conducted a rigid registration experiment with $L=T=1$ and various numbers of steps in ODE solver. Fig. 4(b) shows the corresponding results.

Generalization: We performed the rigid+deformable test on our private dataset. Fig. 4(c) shows bar chart among compared methods.