Enhancing Unsupervised Domain Adaptation via Semantic Similarity Constraint for Medical Image Segmentation

1 Introduction

Medical image segmentation is an important research area, and deep learning is successful in many related tasks with this area. Due to different physical imaging principles and imaging modes, the appearance of achieved digital images will have significant differences, and domain shift is even more severe than natural images. As shown in Figure 1, the magnetic resonance imaging (MRI) and computed tomography (CT) images of the same cardiac region produced markedly different appearances. Well-trained deep neural networks on one domain often fail when being deployed on another domain of the same target. It is necessary to transfer the knowledge learned from one domain to another domain for reducing the performance degradation caused by domain shift. Many efforts have been taken to reduce the performance degradation caused by domain shifting. The unsupervised domain adaptation approach is a solution for it.

∗Corresponding Author.

(b) Examples of CT images

Figure 1: Image examples from CT and MR domains which can illustrate the severe cross-modality domain shift of medical images.

Previous works tackle this problem from three main aspects. One is the image-level adaptation, which reduces domain shift by transforming the images between domains in appearance at pixel-level. These methods commonly adopt the cycle-consistency constraint proposed in CycleGAN [Zhu et al., 2017], which can translate images between unpaired images from different distributions. The second aspect is adapting latent features. The goal of these approaches is to extract domain-invariant features through deep neural networks. Most of these methods are implemented by deploying adversarial training on latent feature distributions. [Tsai et al., 2018] proposed to add discriminator into a more compact space, such as the segmentation mask space. The third aspect is combining the above two perspectives to adapt in image level and feature level simultaneously.

The methods mentioned above perform adaptation via cycle-consistency loss and adversarial learning, which have some disadvantages: 1) if a discriminator is added to a highdimensional and over-complex feature space, it could lead to non-convergence in the GAN training process; 2) domaininvariant semantic feature information could be destructed because of the instability of adversarial training in the segmentation mask space. To solve these problems, we design a novel network structure that can take full advantage of the adversarial learning without destructing the semantic information, and we directly enhance the extracted domain-invariant features via our proposed additional semantic similarity mining module. Our method supports synergistic adaptations from both image and feature spaces simultaneously. The proposed method has two workflows. Specifically, in the first flow, we use a shallow content encoder shared between domains and two domain-specific decoders to translate images between the two domains towards appearance but keep the semantic content. The design of the encoder is shallow because the weight of the encoder is shared by the two modalities. If the content encoder has too much capacity, it would not capture the domain invariant feature but be equivalent to using two independent domain encoders. The image translation workflow is constrained by cycle consistency and adversarial learning. Furthermore, to enhance the extracted domaininvariant feature after the content encoder, we add a strong semantic similarity mining (SSM) task on the content encoder. This task gives the content encoder a strict constraint that the features extracted by the encoder of two augmented images from the same original image should be similar. Another flow uses the synthesis target domain images from the source domain and the original target domain images to train a segmentation model. An adversarial loss is added to the segmentation prediction space between real images and synthesis images. Our experiments show that the proposed method outperforms the state-of-the-art (SOTA) methods with a large margin. The major contributions of this work are listed as below:

• We propose a novel network structure to address the unsupervised domain adaptation problem for medical image segmentation.
• This is the first attempt to add semantic similarity mining to the UDA learning framework to directly enhance the extracted domain-invariant features. The semantic similarity mining task is model-independent and could be easily integrated to other models.
• We evaluate the proposed method on two challenging multi-modal medical structure segmentation datasets, and the proposed method outperforms previous state-ofthe-art methods by a large margin.

2 RelatedWork

Unsupervised domain adaptation is an active research direction recently because it does not require the annotation data of the target domain. It has received significant attention in the field of medical image processing cause heterogeneous domain shift is more severe due to different physical principles. Previous work can be roughly divided into three streams: only feature, only image, image + feature.

1) Feature Alignment. The aim of this stream is to learn domain-invariant semantic features between domains. Most of the models follow the Domain Adversarial Neural Network (DANN) and Adversarial Discriminative Domain Adaptation (ADDA) structure [Ganin et al., 2016; Tzeng et al., 2017]. [Dou et al., 2019] assumed that the severe data shift across domains mainly exists in low-level image features. They proposed to fix the weights of high-level layers in the network and adapt the low-level layers between domains via adversarial learning, and they evaluated the method on the MMWHS dataset [Zhuang and Shen, 2016]. A disadvantage of this method is that it needs to manually select the hyperparameter n, which defines the number of low-level layers.

1. Image Alignment. Driven by the mighty generative power of Generative Adversarial Network [Goodfellow et al., 2014], a lot of works have been proposed in image alignment. [Bousmalis et al., 2017] proposed to translate source domain images to the target domain, then use the target-like images for supervised training, which have inherited labels from the source domain. Inversely, [Zhang et al., 2018b] use the labeled source domain data for supervised training and then transform target domain data to source domain for testing using the pre-trained source model.

3) Image + Feature Alignment. As previously mentioned, the two alignment perspectives are fundamental, and they are complementary to each other in fact. Some recent studies combined these two perspectives to obtain a stronger model for unsupervised domain adaptation [Chen et al., 2020a]. In [Hoffman et al., 2018] and [Zhang et al., 2018a], their methods adapted image-level and feature-level space between synthetic and real-world driving scene domains separately and connected the two stages sequentially, so the internal interactions between the two perspectives can not be achieved. [Chen et al., 2020b] proposed a framework named Synergistic Image and Feature Alignment (SIFA) to conduct synergistic image and feature alignment for cross-modality cardiac image segmentation.

3 Approach

Figure 2 shows the overview of the proposed method. The framework consists of several neural network modules described as follow, a content encoder Ec, two upsampling style decoders {Us, Ut}, a segmentation encoder Es, a pixel classifier C for the final segmentation output, and two MLP blocks.

3.1 Image and Feature Space Alignment

Denote {xs , ys } are samples from the labeled source domain X s and {xt } are samples from the unlabeled target domain X t . There are two appearance translation sub-networks which are Ec ◦Us and Ec ◦Ut who can translate images cross domains or reconstruct images towards appearance. The content encoder Ec is designed not very deep because its weight is shared between two modalities. If the encoder has too much capacity, it would not capture the domain invariant feature but would be equivalent to using two independent domain encoders. The workflows of image translation is shown in Figure 2.

A reconstruction loss is defined to ensure the generated images can preserve the content and style information. The reconstruction loss is defined as:

To ensure the synthesis images xs→t and xt→s to preserve the original semantic content information, we adopt to use the commonly-used cycle-consistency loss [Zhu et al., 2017] for unpaired cross-modality images alignment:

According to CycleGAN [Zhu et al., 2017], the adversarial losses are added on the synthesis images xs→t and xt→s to ensure them like the realistic images:

In above equations, Gs equals to Ec ◦ Ut, Gt equals to Ec ◦ Us. D{s,t} means the corresponding discriminator.

With the constraints of cycle-consistency loss and adversarial loss, the latent feature f = Ec(x{s,t} ) will preserve some semantic content information. At the same time, the distribution of synthesis target domain image xs→t is bringing closer to the real data distribution of the target domain, which can be used to train a target domain segmentation network in a supervised learning manner.

3.2 Segmentation Sub-network

We use a target segmentation encoder Es which accepts input of target-like images to narrow the remaining gap between target domain images xt and synthesis target domain images xs→t . Caused by the instability of GAN training, the generated target-like image xs→t could lose some semantic information. To tackle this problem, we integrate the content encoder Ec into the segmentation sub-module. The xs→t and xt are forwarded into the segmentation encoder Es directly and then the feature extracted by the segmentation encoder Es will be fused with the feature extracted by the content encoder Ec. Finally, a pixel-wise segmentation classifier is applied to predict the final segmentation mask using the information from the content encoder Ec and the segmentation encoder Es. In this part, supervised learning is applied to learn the prediction of segmentation mask using synthesis target-like images xs→t with label ys inherited from xs . The segmentation mask predicted from x{s,t} can be denoted by

We implement the segmentation task learning by minimizing a hybrid loss Lseg:

In above equation, Dice means the widely used Dice loss. We implement Dice loss with Laplacian smoothing, which is defined as below:

The FL termmeans Focal loss [Lin et al., 2017]. It can be regarded as a kind of extension of cross-entropy loss and it is mainly used to tackle the unbalanced samples problem.

where α and γ are hyper-parameters. In our experiments, corresponding values are setting to 0.75 and 2, respectively.

The term T in Equation (5) is the Tversky loss [Salehi et al., 2017]:

where α and β are hyper-parameters, both set to 0.5. Similarly, adversarial learning is applied to ˆ yt and ˆ ys→t .

where S equals to C ◦ (Ec ⊕ Es).

3.3 Additional Constraint via Semantic Similarity Mining

The domain invariant representation is learned from image reconstruction loss, cycle-consistency loss, and adversarial loss. However, these constraints only indirectly affect the domain invariant representation learning. We introduce adding a strong semantic similarity mining (SSM) sub-task into our framework to directly enhance the domain invariance of the feature extracted by the content encoder Ec.

This technique is based on the structure of Siamese networks [Chen and He, 2021]. Siamese networks accept two or more inputs to be processing by weight-sharing networks. The backbone of the Siamese networks is the content encoder Ec. A projection MLP head and a prediction MLP head (denoted as h) are connected to the end of the backbone. We set the encode function f consisting of the backbone and the projection MLP head and we denote z = f(x) and p = h(f(x)) as the outputs of the encode function and the prediction MLP head given the input x. Given {zk i , pki |i = 1, 2 and k = s, t} for the i-th random semantic preserved augmented image xk (k represents domain), we minimize the cosine similarity loss Lsim by the below term:

where || · ||2 is ℓ2-norm, stopgrad is the stop-gradient operation. The model maximizes the cosine-similarity of the outputs between two augmented images from one image.

The existing methods mainly use cycle-consistency and adversarial training to constrain domain-invariant feature learning, both of which are applied at the image level to indirectly guide the feature level alignment. We introduce our proposed semantic similarity mining (SSM) module to directly guide the learning of domain invariance from the deep feature level, which does not require as much data as training a GAN needed. We require the Siamese network to output as similar as possible between two different image-augmented versions of an original image from any modality. The augmentation method of images is well designed for specific modality types to confuse the image appearance between the modalities as much as possible, so we can use this method to mine the semantic similarity cross modalities which helps enhancing domain invariance a lot. Above mentioned random semantic preserved augmentation plays an essential role in the semantic similarity mining task. Every augmentation type in the augmentation set should be semantic preserved, which means the label would not change at all. The augmentation should apply randomly every time.

3.4 Overall Objective

The overall objective function is defined as follows:

The {λrec, λcyc, λadv, λseg, λyadv , λsim} in Equation (12) are trade-off parameters adjusting the weight of each component. The corresponding values are set to {10.0, 1.0, 1.0, 1.0, 0.01, 1.0} respectively. The Lrec, Lcyc, Ladv, and Lseg are similar to previous works. Qualitatively, it can be considered that the following three types of constraints are successively weakened: the reconstruction constraints Lrec > the domain alignment constraints Lcyc and Ladv > the segmentation realistic constraints Lyadv . So these hyper-parameters are set to 10,1,0.01 respectively. The last term Lsim is specific to our method, it is a loss term that also constrains the alignment between modalities, so it is set to 1, too. We tried other values like 10, 0.1, and our method performs best in the value of 1.

3.5 Optimization Steps

The entire training process is End-2-End, we do not train every module separately. But in each training iteration, we first optimize the image translation module and semantic similarity mining module together, then optimize the discriminators of source and target domains. After that, we update the parameters of the segmentation module. More training details and network structures could be found in the supplementary material.

4 Experiments

4.1 Dataset

Cardiac substructure segmentation. Multi-Modality Whole Heart Segmentation Challenge 2017 dataset (MMWHS) [Zhuang and Shen, 2016] is a cardiac segmentation dataset including two modality images (MR and CT). Each modality contains 20 volumes collected from different sites, and there is no pair relationship between modalities. We choose four classes of the cardiac structures. They are the ascending aorta (AA), the left atrium blood cavity (LAC), the left ventricle blood cavity (LVC), and the myocardium of the left ventricle (MYO). Abdominal multi-organ segmentation. We choose the public CT data from [Landman et al., 2017] (30 volumes) and T2-SPIR MRI training data from the ISBI 2019 CHAOS Challenge [Selver et al., 2019] (20 volumes) to validate our method. Ground truth segmentation masks of four abdominal organs including liver (L), right kidney (R.K), left kidney (L.K), and spleen (SP) are provided in both modalities.

For each dataset, we randomly split them with 70% cases for training, 10% cases for validation and 20% cases for test. For training our model, online augmentation is applied to reduce over-fitting, which mainly includes 2D/3D random elastic deformation, random rotation, random scaling, and 3D affine transformation. The final training data is 2D coronal slices in the resolution of 256x256x1 sampled from the augmented volumes.

Abdominal MR to Abdominal CT

4.2 EvaluationMetrics

We choose the Dice similarity coefficient (Dice) and the average surface distance (ASD) as our evaluation metrics. Dice similarity coefficient is a measurement that measures the volume overlap between the prediction and corresponding ground truth. Average surface distance evaluates the distance between the surface of the prediction and the label. A higher Dice or a lower ASD means better performance.

4.3 Comparison with Other Methods

In this section, we compared our method with several stateof-the-art methods in quantitative and qualitative. PnPAdaNet [Dou et al., 2019] proposed a plug-and-play style method to align the feature spaces of the domains; AdaOutput [Tsai et al., 2018] conducts an adversarial task between multi-level prediction maps; CycleGAN [Zhu et al., 2017] performs an image-to-image translation by a cyclist reconstruction manner; CyCADA [Hoffman et al., 2018] aligns the feature from different domains at both image and feature level; SIFA [Chen et al., 2020b] conducts synergistic alignment of the domain from both image and feature perspectives; DSFN [Zou et al., 2021] introduces a dual-scheme network structure to reduce the domain gap; DSAN [Han et al., 2021] proposed a deep symmetric architecture to eliminate the domain gap. In the above-mentioned methods, AdaOutput, CycleGAN, and CyCADA are proposed for natural image datasets, but PnP-AdaNet, SIFA, DSFN and DSAN are proposed for medical images.

Table 1 shows the comparison on the two datasets in the adaptation direction of MR to CT. Supervised training result for this task is took as the performance upper bound, and the result that the model trained on the source domain but directly applied to the target domain is obtained as the lower bound. Both of them are recorded in the "Supervised training W/o adaptation". In the adaptation direction of cardiac MR to cardiac CT, according to the "Supervised training W/o adaptation" result, we can obtain that the apparent differences caused by domain shift are severe, which drop the performance of average 90.9% to 17.2% in Dice coefficient. With our proposed unsupervised domain adaptation method, we achieve the average Dice coefficient of 80.8% and the ASD value of 4.2 remarkably. In the abdominal dataset, result is similar but the task is more easy. Visual comparison results are provided in Figure 3. We can see that in the "W/o adaptation" case, the prediction is definitely wrong. Our method outperforms all of the other methods. More experiment results are reported in supplementary material.

4.4 Ablation Study

We conducted a number of ablation studies to evaluate the effectiveness of each key component in our proposed model. The results are shown in Table 2.

The baseline method is using image adaptation only, this method train the image translation-related modules, and test on an extra segmentation network using target-like images. Compared to "W/o adaptation" lower-bound, the result shows that this method effectively narrows the gap between domains because the Dice increased to 58.5% compared to the "W/o adaptation" case.

Next, we study the necessity of sharing weights between encoders. We conduct a study that using two independent content encoders for every domain. In this case, the average Dice dropped 9.4%. It proves that sharing the weights of the content encoder between domains is necessary for extracting domain-invariant features.

We use a light-weight nine layers convolutional neural network as the content encoder, which is relatively shallower than the commonly used deep structures which usually have dozens of layers, and we also conduct experiments that change the number of layers in the encoder Ec. The average Dice comparison of different number of layers in the encoder Ec is shown in Table 3.

The semantic similarity loss is an important part of our method, and we also study its effectiveness. In comparison, we remove this loss term in the overall loss term and train again. We can see a reduction of 3.9% in the average Dice compared to the original model.

5 Conclusion

We proposed a novel unsupervised domain adaptation medical image segmentation algorithm that combines both unpaired image transformation and domain invariant feature alignment learning. Domain invariant features are directly enhanced by the proposed SSM module. Extensive experiments with promising results on two challenging datasets show the effectiveness of our approach. At present, all of the experiments only use 2D volume slices, and it is also possible to extend our method to 3D volumes. From a practical perspective, our method has the potential to be extended to evaluate the quality of cross-domain image translation, e.g., via scoring the segmentation predictions of real and translated images by the discriminator. These could be explored in our future work.

Acknowledgments

This work was supported by the NSFC Projects 62076096 and 62006078, the Shanghai Municipal Project 20511100900, Shanghai Knowledge Service Platform Project ZF1213, STCSM Project 22ZR1421700 and the Fundamental Research Funds for the Central Universities.