Unsupervised Multi-Modal Medical Image Registration via Discriminator-Free Image-to-Image Translation

1 Introduction

Different medical image modalities, such as Magnetic Resonance Imaging (MRI) and Computed Tomography (CT), show unique tissue features due to the acquisition with different scanners and protocols. Multiple image modalities can be fused to provide combined information, which is known as the process of multi-modal image fusion [Pielawski et al., 2020]. A wide variety of clinical applications rely on the fusion of multi-modal data, e.g., preoperative planning and image-guided radiotherapy [Oh and Kim, 2017]. Since the images need to be aligned first through registration for image fusion, it is of great importance to establish anatomical correspondences among images of different modalities by using multi-modal image registration.

∗Corresponding Author

Learning-based registration seeks to predict deformation fields directly from a pair of images by maximizing a predefined similarity metric [Fan et al., 2019]. Supervised or semisupervised learning strategies use ground-truth deformation fields or segmentation masks in the training phase, and may suffer from the lack of data labeling [Uzunova et al., 2017; Hu et al., 2018]. Since it is extremely time-consuming and laborious to label registration data even for specialists, unsupervised methods have been proposed to overcome this limitation solely by maximizing the image similarity between the target image and the source image. However, the performance of unsupervised methods is highly dependent on the choice of cross-modal similarity metrics. Generally, widespread similarity metrics like the sum of squared differences (SSD) and normalized cross correlation (NCC), which are well-suited for mono-modal registration problems [Balakrishnan et al., 2019; de Vos et al., 2017], perform badly in a multi-modal setting. Typically, unsupervised multi-modal registration approaches use Normalized Mutual Information (NMI) and Modality-Independent Neighbourhood Descriptor (MIND) [Maes et al., 2003; Heinrich et al., 2012]. Since NMI, as a global metric, only measures statistical dependence between two entire images, it is difficult to use it for local image alignment. MIND, on the other hand, is a patch-based image similarity metric, which tends to suffer from severe image deformations and cannot achieve global alignment.

Given the recent success of multi-modal image translation [Huang et al., 2018; Park et al., 2020], an alternative solution for addressing multi-modal registration is to convert the problem to a simpler unimodal task using an image-toimage (I2I) translation framework [Qin et al., 2019]. Specifically, translation-based methods use Generative Adversarial Network (GAN) mode to translate images from source modality to target modality. And the GAN consists of a generator and a discriminator, where the generator learns to generate plausible data and the discriminator penalizes the generator for producing unrealistic results. With the GAN mode, the registration network can be trained with unimodal similarity metrics. However, this GAN-based image translation tends to produce shape inconsistency and artificial anatomical features which will, in turn, deteriorate the performance of the registration [Arar et al., 2020; Xu et al., 2020]. More specifically, different modalities have very distinct geometric variances caused by the shape of the imaging bed, the imaging protocols of the scanner, and the field of view. We refer to these variances as "domain-specific deformations" [Wang et al., 2021]. We argue that the inconsistency and artifacts are introduced by the discriminator that mistakenly encodes domain-specific deformations as indispensable appearance features and encourages the generator to reproduce the deformations. This tends to create unnecessary difficulty for registration tasks. This paper shows that we can improve the performance of multi-modal image registration by removing the discriminator in image-to-image translation.

In this work, we propose a novel translation-based unsupervised registration approach for aligning multi-modal images. Our main idea is to reduce the inconsistency and artifacts of the translation by removing the discriminator as we discussed above. Specifically, we replace the GAN-based translation network with our discriminator-free translation network (shown in Figure 1). The presented translation network incorporates a patchwise contrastive loss (PatchNCE [Park et al., 2020]) to maintain shape consistency during translation and a pixel loss to integrate the output with the target appearance. Additionally, two novel loss terms are also proposed, one is local alignment loss and the other is global alignment loss. The local alignment loss captures detailed local texture information by modeling the detail-rich image patches. The global alignment loss focuses on the overall shape for the detail-missing generated images. The proposed translationbased registration method coupled with these two losses can achieve local and global alignment and yield more accurate deformation fields.

The main contributions of our work can be summarized as follows:

• We present a discriminator-free I2I translation mode to replace the original GAN-based I2I mode and achieve accurate registration.
• We design a contrastive PatchNCE loss in our translation-based registration model as a shapepreserving constraint.
• We also propose a local and a global alignment losses to further improve the registration performance.

2 RelatedWork

In recent years, many translation-based multi-modal registration approaches have been proposed. They commonly follow a registration-by-translation framework: an I2I translation network is first trained to synthesize fake images with the target appearance, and then mono-modality metrics can be used in the target domain. In unsupervised translation, cycleconsistent generative adversarial networks (CycleGAN [Zhu et al., 2017]) are widely adopted. However, cycle consistency leads to multiple solutions, which means that the translated images can not maintain the anatomical structure of source images and may contain artifacts [Kong et al., 2021]. With inaccurate image translation, the performance of multi-modal registration tends to be degraded. To address this difficulty, other approaches beyond CycleGAN have been proposed.

Qin et al. [2019] use image disentanglement to decompose images into common domain-invariant latent shape features and domain-specific appearance features. Then the latent shape features of both modalities are used to train a registration network. But this method still relies on cycle consistency and a GAN mode, which inevitably introduces inconsistency and hampers the process of registration.

Arar et al. [2020] attempt to force the translation and the registration steps to be commutative, which can implicitly encourage the translation network to be structure-preserving. The structure-preserving translation network allows the use of simple mono-modality metrics for training a registration network. However, their method is still GAN-based, which means the structure consistency will be affected by the presence of the discriminator.

Closest to our work, Casamitjana et al. [2021] propose a synthesis-by-registration method, which is different from the previous registration-by-synthesis methods. Their approach is made up of two stages: stage one is training a registration network on pairs of images from the target domain with data augmentation, and stage two is training an I2I network while freezing the parameters of the registration network. However, their registration network is first trained on the images from the target modality instead of images from the two modalities, which may guide the registration network to generate an unrealistic deformation field. And both the registration network and translation network are used in test time in their scheme. On the contrary, our method is a joint framework with end-toend optimization and only the registration network is needed in test time, which leads to a more reliable deformation field.

3 Method

In this work, we propose an end-to-end learning framework for registering multi-modal image pairs in a fully unsupervised manner. Our core idea is to replace GAN-based I2I translation with our novel discriminator-free I2I translation, shown in Figure 1. The PatchNCE loss maintains shape consistency and the pixel loss enables the appearance transfer. With the proposed discriminator-free I2I translation, the multi-modal registration task is converted into a unimodal one.

Our method consists of two components: (1) a deformable registration network, and (2) a discriminator-free translation network for mapping images from source domain to target domain and reconstructing images from the target domain. The two components are trained jointly, and only the registration network is used in test time. The pipeline of our method is depicted in Figure 2(a).

We denote the registration network as R and the translation network as T, as in Figure 2. Let X and Y denote two paired image domains, where X is source domain and Y is target domain. Pairing means that each image x ∈ X has a corresponding unaligned image y ∈ Y representing the same anatomical structure. The process of registration is to find a deformation field that aligns the source image x to the target image y accurately. Given an image pair (x, y) as input, R learns to predict a deformation field φ, which describes how to non-rigidly align x to y. Meanwhile, T takes x and y as the inputs and outputs target-modality images y′ and ˆ y, where y′ = T(x) and ˆ y = T(y). y′ is the translated image, which has similar appearance as images in domain Y. ˆ y is the reconstructed image of y. The PatchNCE loss LPatchNCE is employed to force y′ and ˆ y to keep the orginial structure of x and y, respectively. x(φ) and y′ (φ) are warped images from the source image x and the translated image y′ . The pixel loss Lappearance enables the appearance transfer of network T between y′ (φ) and y. In addition, local alignment loss Llocal and global alignment loss Lglobal are employed to further improve the registration performance.

3.1 Registration Network

The registration network R takes an image pair (x, y) as an input and outputs a deformation field φ = R(x, y). The warped image x(φ) is aligned with y. In a two-dimensional setting, the deformation field is a matrix of 2D vectors, indicating the moving direction for every pixel in the source image x. To generate smooth deformation fields and penalize the tendency of overly distorting the deformed image x(φ), we adopt an L2-norm of the gradients of the deformation field as the regularization term [Hoopes et al., 2021], which is denoted as Lsmooth. Formally, the loss at pixel v = (i, j) is given by:

where N(v) denotes a set of neighbor pixels of the v.

3.2 Discriminator-Free Translation Network

Our translation network T takes images from the source domain X and outputs translated images that have similar appearance as images in the target domain Y. We divide our translation network T into two components, an encoder Tenc followed by a decoder Tdec, shown in Figure 2(b). Tenc extracts shape-related features, while Tdec learns to perform shape-preserving modality translation with those features. Given the input x, Tenc and Tdec jointly generate the output y′ = T (x) = Tdec (Tenc(x)).

A key task of our method is to train the I2I translation network T to translate images without the discriminator. If the output from T is shape-preserving, which implies T cannot deform the anatomical structure, the alignment task will be done solely by registration network R. Therefore, we propose the contrastive PatchNCE loss to enforce the shape consistency and the pixel loss to enable the appearance transfer from the source modality to the target modality.

PatchNCE Loss. The PatchNCE loss maximizes the mutual information between input image patches and output image patches, which is based on a noise contrastive estimation framework [Oord et al., 2018]. It makes every output patch similar to the corresponding input patch, while different from the other patches within the input. We use a "query" to refer to an output patch and "positive" for the corresponding input patch and "negatives" for the noncorresponding input patches, shown in Figure 2(b). For similarity measurement, the encoder part Tenc of our translation network with additional two-layer multi-layer perceptron (MLP) is used to map image patches to embedded vectors. Specifically, the query, positive, and N negatives are embedded to K-dimensional vectors z, z+ ∈ RK and z- ∈ RN×K , respectively. zn ∈ RK denotes the n-th negative patch. We convert the similarity measurement to an (N+1)-way classification problem, where the similarity between the query and other samples are expressed as logits. The cross-entropy loss for multi-class classification is calculated, representing the probability of the positive being selected over N negatives. The formula is given by:

where τ is the temperature parameter set to 0.07 in our experiments.

Given an input image, the encoder part Tenc will generate multilayer hidden features, which form a feature stack. A spatial location in a layer of the feature stack represents a patch of the input image, with deeper layers corresponding to bigger patches. Let L denote the number of extracted layers from the feature stack. A MLP Hl with two layers is used to map the selected encoder features to embedded representations {zl}L = {Hl (T l enc (x))}L, where T l enc represents the features of the l-th selected layer and l ∈ {1, 2, . . . ,L}. Let s ∈ {1, . . . , Sl}, where Sl represents the number of spatial locations in each T l enc . For each spatial location s, the corresponding embedded code is referred to as zs l ∈ RK , and the features at any other locations are denoted by zS\s l ∈ R(Sl-1)×K . Similarly, the corresponding embedded representations of output image y′ are {z′ l }L = {Hl (T l enc (G(x)))}L.

With features from multiple layers of the encoder, patchwise noise contrastive estimation can be applied on multiple scales. Multilayer PatchNCE loss is given by:

In addition, LPatchNCE is also used on images from target domain Y, which acts as reconstruction loss. Through the loss LPatchNCE(T, H, Y ), network T outputs the reconstruction ˆ y = T (y).

Pixel Loss. Contrastive loss is effective to preserve the shape of the input image x. However, without adversarial loss, we still need to maximize the appearance similarity between a translated image y′ and the target domain Y with a pixel loss. We define the pixel loss with L1-norm as:

where y′ (φ) indicates the warped image of y′ .

Lappearance explicitly penalize the absolute intensity differences between y′ (φ) and y′ . The combination of Lappearance and LPatchNCE leads to a discriminaotr-free and shape-preserving translation. Note that by minimizing the two losses, the registration network R is trained jointly to predict a deformation field φ, which aligns y′ to y.

3.3 Local and Global Alignment

To enable R to learn the alignment at the local (patch) level, we propose a variant of the PatchNCE loss. Similarly, we refer to a patch of the warped source image x(φ) as a "query", while "positive" and "negatives" are corresponding and noncorresponding patch(es) within the target image y, respectively. The warped source image and the target image are mapped to embedded vectors {ql}L = {Hl (T l enc (x(φ)))}L and {pl}L = {Hl (T l enc (y))}L, respectively. The patchwise noise contrastive estimation is computed between the embedded vectors {ql}L and {pl}L, which is different from the original PatchNCE. For clarity, we denote this variant loss as Llocal:

Applied on the cross-modality image patches, Llocal encourages the registration network R to learn local alignment as shown in Figure 3. Note that the images produced by our discriminator-free translation network do not contain image texture information. This is beneficial to the extraction of global information, shown in Figure 3. Inspired by this, we further propose a global alignment loss:

Minimizing Lglobal leads to similar style between the generated y′ and ˆ y. Meanwhile, the registration network R learns a deformation field φ to best align y′ to ˆ y.

3.4 Final Objective

Our final objective is as follows:

where we set λP = 0.25, λA = 1, λL = 0.25 and λG = 1 in our experiments.

4 Experiments

4.1 Datasets

We evaluated our proposed method on two public datasets. Both of them are obtained from MICCAI Learn2Reg 2021 challenge [Hering et al., 2021], which is a comprehensive registration challenge covering different anatomical structures and modalities. Specifically, the first dataset is for the task of CT-MR thorax-abdomen intra-patient registration, and the other one is for the task of CT lung inspiration-expiration registration.

Thorax-Abdomen CT-MR Dataset. This dataset contains 16 pairs of CT and MR abdomen scans. The annotations on all scans are manual and automatic segmentations of multiple organs. Each scan is a 3D volume in size of 192×160×192 with 2 mm voxel spacing. After preprocessing the data, e.g., coarse affine registration using the Elastix toolbox [Marstal et al., 2016], we randomly split the dataset into 10/2/4 pairs for train/validation/test, and central 90 slices containing organs in each scan are extracted for our experiments. All slices are padded and resized into 256 × 256.

Lung CT Inspiration-Expiration (Insp-Exp) Dataset. The dataset consists of 30 pairs of CT lung scans. All lungs are segmented automatically or manually. The challenge in this dataset is to estimate the underlying breathing motion between inspiration and expiration. Each scan is a 3D volume in size of 192 × 192 × 208 with resolution of 1.75 mm. In our experiments, we randomly divide the 30 pairs into 20/4/6 for train/validation/test, and on each volume, we extract middle 100 slices. And each slice is resized into 256 × 256. To simulate a multi-modal registration task with this singlemodal dataset, we synthesize a new modality using the intensity transformation function cos (I · π/255), Gaussian blurring with the kernel size of 3 × 3 and random elastic deformation sequentially as proposed in [Qin et al., 2019]. Note that the intensity transformation function is only applied to the foreground of each slice. Deformation fields are predicted between real CT inspiration slices and their corresponding synthesized expiration slices.

4.2 Implementation Details

We implement our model based on the framework and implementation of CUT [Park et al., 2020]. The translation network T is a Resnet-based architecture with 9 residual blocks. Our encoder is defined as the first half of the translation network, and five layers of features in the encoder are extracted. The registration network adopts a U-net based architecture with skip connections from contracting path to expanding path [Ronneberger et al., 2015]. For the initialization of networks, we use the Xavier initialization method. Our networks are implemented in PyTorch and all the experiments were conducted on GeForce RTX 2080 Ti. We use Adam optimizer to train our model for 300 epochs with parameters lr = 0.0002, β1 = 0.5 and β2 = 0.999. Linear learning rate decay is activated after 200 epochs.

4.3 Evaluation

Metrics. For the thorax-abdomen dataset, we directly use multi-organ segmentation masks to evaluate the registration accuracy. Dice similarity coefficient (DSC [Dice, 1945]) and 95% percentile of Hausdorff distance (HD95 [Huttenlocher et al., 1993]) are computed between multi-organ masks of a warped source image and its target image. Similarly, DSC and HD95 are calculated between the provided lung masks for the evaluation of the lung CT dataset. DSC is used to measure the accuracy of registration, while HD95 measures reliability. A higher DSC and lower HD95 indicate a better performance of the registration model.

Baselines. We compare our method against several recent state-of-the-art multi-modal registration methods and some other well-established methods. Specifically, the competing methods are: (1) Affine, affine registration based on the normalized mutual information using the Elastix toolbox; (2) MIND, a VoxelMorph architecture [Balakrishnan et al., 2019] with similarity metric MIND; (3) CGAN, a CycleGAN, which is pre-trained on unpaired images, combines with the VoxelMorph registration network using mono-modal similarity metric NCC; (4) RGPT, a multi-modal registration model via geometry preserving translation [Arar et al., 2020]; (5) SbR, a recent synthesis-by-registration model based on contrastive learning [Casamitjana et al., 2021]. For ablation study, we propose three variants of our model by removing the loss terms one by one: ours w/o Llocal, ours w/o Lglobal, and ours w/o Llocal and Lglobal (ours w/o l & g). To investigate if the inaccuracy is introduced by the discriminator, we also study a variant with a discriminator plugged into the translation network in our model (ours w/ dis).

4.4 Results

The quantitative results on the two datasets are summarized in Table 1 and Table 2. We compare the proposed method with the other five methods by measuring the registration accuracy with the DSC and HD95. Our method consistently outperforms other competing methods. Our registration network can predict more accurate deformation fields, even when there exists significant shape deformation and style difference between source images and target images. In addition, Figure 4 displays four examples of the warped source images and their corresponding deformation fields. The deformation fields generated by the RGPT method are corrupted by noise and jitter (the red boxes in column 5-6 in Figure 4). This explains its inferior performance in HD95. Even though RGPT can achieve a good boundary alignment, the image quality is degraded after the registration. On the contrary, SbR produces relatively smooth deformation fields but fails to deal with large-scale deformations (the red boxes in column 7-8 in Figure 4). This is because the registration network of SbR is pre-trained on the images from the target modality with randomly generated deformation fields and its parameters are frozen once the pre-training is done. Different from the above methods, ours achieves the most accurate boundary alignment by capturing both local and global information without any discriminator. (the red boxes in column 3-4 in Figure 4).

4.5 Ablation Study

The primary objective of our work is to achieve accurate multi-modal registration. Therefore, we design three variants of the proposed model to investigate the impact of Llocal and Lglobal. The results in Table 1 and 2 show that, Llocal and Lglobal are complementary, since the proposed method outperform the three ablated versions. What's more, with only one of the two losses, the registration performance still can be improved. To investigate whether the discriminator will introduce inconsistency and degrade the performance of the registration, we design a variant, ours w/ dis. As can be seen in the last row of Table 1 and 2, the registration performance becomes worse than the proposed method.

4.6 Conclusion

We propose a novel discriminator-free and shape-preserving translation network for multi-modal registration, taking advantage of contrastive learning. The registration network successfully gets rid of inconsistency and artifacts introduced by the discriminator. The contrastive loss ensures shape consistency while the pixel loss enables the appearance transfer. Furthermore, we leverage a local and a global alignment losses to achieve local and global alignment, improving the registration accuracy. Finally, we evaluate the proposed method on two open datasets, and show that it outperforms the state-of-the-art methods.

Acknowledgments

This work is supported in part by the Natural Science Foundation of Guangdong Province (2020A1515010717), NSF1850492 and NSF-2045804.