Accurate delineation of lesions or anatomical structures is a vital step for clinical diagnosis, intervention, and treatment planning [tajbakhsh2020embracing]
. While recently flourished deep learning methods excel at segmenting those structures, deep learning-based segmentors cannot generalize well in a heterogeneous domain, e.g., different clinical centers, scanner vendors, or imaging modalities[liu2022deep, Liu_2021_ICCV, liuconstraining, che2019deep]. To alleviate this issue, unsupervised domain adaptation (UDA) has been actively developed, by applying a well-performed model in an unlabeled target domain via supervision of a labeled source domain [chen2019synergistic, liu2021domain, liu2021generative, liu2022unsupervisedFrontiers]. Due to diverse target domains, however, the performance of UDA is far from satisfactory [zou2020unsupervised, han2022deep, liu2022self]. Instead, labeling a small set of target domain data is usually more feasible [van2020survey]. As such, semi-supervised domain adaptation (SSDA) has shown great potential as a solution to domain shifts, as it can utilize both labeled source and target data, in addition to unlabeled target data. To date, while several SSDA classification methods have been proposed [donahue2013semi, yao2015semi, saito2019semi, kim2020attract], based on discriminative class boundaries, they cannot be directly applied to segmentation, since segmentation involves complex and dense pixel-wise predictions.
Recently, while a few works [wang2020alleviating, chen2021semi, hoyer2021improving]
have been proposed to extend SSDA for segmentation on natural images, to our knowledge, no SSDA for medical image segmentation has yet been explored. For example, a depth estimation for natural images is used as an auxiliary task as in[hoyer2021improving], but that approach cannot be applied to medical imaging data, e.g., MRI, as they do not have perspective depth maps. Wang et al. [wang2020alleviating] simply added supervision from labeled target samples to conventional adversarial UDA. Chen et al. [chen2021semi] averaged labeled source and target domain images at both region and sample levels to mitigate the domain gap. However, source domain supervision can easily dominate the training, when we directly combine the labeled source data with the target data [saito2019semi]. In other words, the extra small amount of labeled target data has not been effectively utilized, because the volume of labeled source data is much larger than labeled target data, and there is significant divergence across domains [saito2019semi].
To mitigate the aforementioned limitations, we propose a practical asymmetric co-training (ACT) framework to take each subset of data in SSDA in a unified and balanced manner. In order to prevent a segmentor, jointly trained by both domains, from being dominated by the source data only, we adopt a divide-and-conquer strategy to decouple the label supervisions for the two asymmetric segmentors, which share the same objective of carrying out a decent segmentation performance for the unlabeled data. By “asymmetric,” we mean that the two segmentors are assigned different roles to utilize the labeled data in either source or target domain, thereby providing a complementary view for the unlabeled data. That is, the first segmentor learns on the labeled source domain data and unlabeled target domain data as a conventional UDA task, while the other segmentor learns on the labeled and unlabeled target domain data as a semi-supervised learning (SSL) task. To integrate these two asymmetric branches, we extend the idea of co-training [blum1998combining, balcan2005co, qiao2018deep], which is one of the most established multi-view learning methods. Instead of modeling two views on the same set of data with different feature extractors or adversarial sample generation in conventional co-training [blum1998combining, balcan2005co, qiao2018deep], our two cross-domain views are explicitly provided by the segmentors with the correlated and complementary UDA and SSL tasks. Specifically, we construct the pseudo label of the unlabeled target sample based on the pixel-wise confident predictions of the other segmentor. Then, the segmentors are trained on the pseudo labeled data iteratively with an exponential MixUp decay (EMD) scheme for smooth propagation. Finally, the target segmentor carries out the target domain segmentation.
The contributions of this work can be summarized as follows:
We present a novel SSDA segmentation framework to exploit the different supervisions with the correlated and complementary asymmetric UDA and SSL sub-tasks, following a divide-and-conquer strategy. The knowledge is then integrated with confidence-aware pseudo-label based co-training.
An EMD scheme is further proposed to mitigate the noisy pseudo label in early epochs of training for smooth propagation.
To our knowledge, this is the first attempt at investigating SSDA for medical image segmentation. Comprehensive evaluations on cross-modality brain tumor (i.e., T2-weighted MRI to T1-weighted/T1ce/FLAIR MRI) segmentation tasks using the BraTS18 database demonstrate superiority performance over conventional source-relaxed/source-based UDA methods.
In our SSDA setting for segmentation, we are given a labeled source set , a labeled target set , and an unlabeled target set , where , , and are the number of samples for each set, respectively. Note that the slice , and , and the segmentation mask labels , and have the same spatial size of . In addition, for each pixel or indexed by , the label has classes, i.e., . There is a distribution divergence between source domain samples, , and target domain samples, and . Usually, is much smaller than . The learning objective is to perform well in the target domain.
2.1 Asymmetric Co-training for SSDA segmentation
To decouple SSDA via a divide-and-conquer strategy, we integrate with either or to form the correlated and complementary sub-tasks of UDA and SSL. We configure a cross-domain UDA segmentor and a target domain SSL segmentor , which share the same objective of achieving a decent segmentation performance in . The knowledge learned from the two segmentors is then integrated with ACT. The overall framework of this work is shown in Fig. 1.
Conventional co-training has focused on two independent views of the source and target data or generated artificial multi-views with adversarial examples, which learns two classifiers for each of the views and teaches each other on the unlabeled data[blum1998combining, qiao2018deep]. By contrast, in SSDA, without multiple views of the data, we propose to leverage the distinct yet correlated supervision, based on the inherent discrepancy of the labeled source and target data. We note that the sub-tasks and datasets adopted are different for the UDA and SSL branches. Therefore, all of the data subsets can be exploited, following well-established UDA and SSL solutions without interfering with each other.
To achieve co-training, we adopt a simple deep pseudo labeling method [wei2020theoretical], which assigns the pixel-wise pseudo label for . Though UDA and SSL can be achieved by different advanced algorithms, deep pseudo labeling can be applied to either UDA [zou2019confidence] or SSL [wei2020theoretical]. Therefore, we can apply the same algorithm to the two sub-tasks, thereby greatly simplifying our overall framework. We note that while a few methods [xia2021uncertainty] can be applied to either SSL or UDA like pseudo labeling, they have not been jointly adopted in the context of SSDA.
Specifically, we assign the pseudo label for each pixel in with the prediction of either or , therefore constructing the pseudo labeled sets and for the training of another segmentor and , respectively:
are the predicted probability of classw.r.t. using and , respectively. is a confidence threshold. Note that the low softmax prediction probability indicates the low confidence for training [zou2019confidence, liu2021generative]. Then, the pixels in the selected pseudo label sets are merged with the labeled data to construct and for the training of and with a conventional supervised segmentation loss, respectively. Therefore, the two segmentors with asymmetrical tasks act as teacher and student of each other to distillate the knowledge with highly confident predictions.
2.2 Pseudo-label with Exponential MixUp Decay
Initially generated pseudo labels with the two segmentors are typically noisy, which is significantly acute in the initial epochs, thus leading to a deviated solution with propagated errors. Numerous conventional co-training methods relied on simple assumptions that there is no domain shift, and the predictions of the teacher model can be reliable and be simply used as ground truth. Due to the domain shift, however, the prediction of in the target domain could be noisy and lead to an aleatoric uncertainty [der2009aleatory, kendall2017uncertainties, hu2019supervised]. In addition, insufficient labeled target domain data can lead to an epistemic uncertainty related to the model parameters [der2009aleatory, kendall2017uncertainties, hu2019supervised].
To smoothly exploit the pseudo labels, we propose to adjust the contribution of the supervision signals from both labels and pseudo labels as the training progresses. Previously, vanilla MixUp [zhang2017mixup] was developed for efficient data augmentation, by combining both samples and their labels to generate new data for training. We note that the MixUp used in SSL [berthelot2019mixmatch, chen2021semi] adopted a constant sampling, and did not take the decay scheme for gradual co-training. Thus, we propose to gradually exploit the pseudo label by mixing up or with pseudo labeled , and adjust their ratio with the EMD scheme. For the selected and with the number of slices and , we mix up each pseudo labeled image with all images from or to form the mixed pseudo labeled sets and . Specifically, our EMD can be formulated as:
where is the MixUp parameter with the exponential decay w.r.t. iteration . is the initial weight of ground truth samples and labels, which is empirically set to 1. Therefore, along with the increase over iteration , we have smaller , which adjusts the contribution of the ground truth label to be large at the start of the training, while utilizing the pseudo labels at the later training epochs. Therefore, and gradually represent the pseudo label sets of and . We note that the mixup operates on the image level, which is indicated by . The number of generated mixed samples depends on the scale of and in each iteration and batch size . With the labeled , , as well as the pseudo labeled sets with EMD and , we update the parameters of the segmentors and , i.e., and with SGD as:
where indicates the learning rate, and denotes the learning loss on with the current segmentor parameterized by . The training procedure is detailed in Algorithm 1. After training, only the target domain specific SSL segmentor is used for testing.
3 Experiments and Results
To demonstrate the effectiveness of our proposed SSDA method, we evaluated our method on T2-weighted MRI to T1-weighted/T1ce/FLAIR MRI brain tumor segmentation using the BraTS2018 database [menze2014multimodal]. We denote our proposed method as ACT, and used ACT-EMD for an ablation study of an EMD-based pseudo label exploration.
Of note, the BraTS2018 database contains a total of 285 patients [menze2014multimodal] with the MRI scannings, including T1-weighted (T1), T1-contrast enhanced (T1ce), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (FLAIR) MRI. For the segmentation labels, each pixel belongs to one of four classes, i.e., enhancing tumor (EnhT), peritumoral edema (ED), necrotic and non-enhancing tumor core (CoreT), and background. In addition, the whole tumor covers CoreT, EnhT, and ED. We follow the conventional cross-modality UDA (i.e., T2-weighted to T1-weighted/T1ce/FLAIR) evaluation protocols [zou2020unsupervised, han2022deep, liu2022self] for 8/2 splitting for training/testing, and extend it to our SSDA task, by accessing the labels of 1-5 target domain subjects at the adaptation training stage. All of the data were used in a subject-independent and unpaired manner. We used SSDA:1 or SSDA:5 to denote that one or five target domain subjects are labeled in training.
For a fair comparison, we used the same segmentor backbone as in DSA [han2022deep] and SSCA [liu2022self], which is based on Deeplab-ResNet50. Without loss of generality, we simply adopted the cross-entropy loss as , and set the learning rate and confidence threshold . Both and
have the same network structure. For the evaluation metrics, we adopted the widely used DSC (the higher, the better) and Hausdorff distance (HD: the lower, the better) as in[han2022deep, liu2022self]
. The standard deviation was reported over five runs.
|Source Only||No DA||4.2||65.2||6.3||27.71.2||55.7||28.0||49.8||39.60.5|
|DICE||Score (DSC)||||Hausdorff||Distance (HD)||[mm]|
|Source Only||No DA||20.61.0||39.50.8||41.30.9||54.70.4||55.20.6||42.50.4|
The quantitative evaluation results of the whole tumor segmentation are provided in Table 1. We can see that SSDA largely improved the performance over the compared UDA methods [han2022deep, liu2022self]. For the T2-weighted to T1-weighted MRI transfer task, we were able to achieve more than 10% improvements over [han2022deep, liu2022self] with only one labeled target sample. Recent SSDA methods for natural image segmentation [wang2020alleviating, chen2021semi] did not take the balance between the two labeled supervisions into consideration, easily resulting in a source domain-biased solution in case of limited labeled target domain data, and thus did not perform well on target domain data [saito2019semi]. In addition, the depth estimation in [hoyer2021improving] cannot be applied to the MRI data. Thus, we reimplemented the aforementioned methods [wang2020alleviating, chen2021semi] with the same backbone for comparisons, which is also the first attempt at the medical image segmentation. Our ACT outperformed [wang2020alleviating, chen2021semi] by a DSC of 3.3% w.r.t. the averaged whole tumor segmentation in SSDA:1 task. The better performance of ACT over ACT-EMD demonstrated the effectiveness of our EMD scheme for smooth adaptation with pseudo-label. We note that we did not manage to outperform the supervised joint training, which accesses all of the target domain labels, which can be considered an “upper bound” of UDA and SSDA. Therefore, it is encouraging that our ACT can approach joint training with five labeled target subjects. In addition, the performance was stable for the setting of from 1 to 10.
In Table 2, we provide the detailed comparisons for more fine-grained segmentation w.r.t. CoreT, EnhT, and ED. The improvements were consistent with the whole tumor segmentation. The qualitative results of three target modalities in Fig. 2 show the superior performance of our framework, compared with the comparison methods.
In Fig. 3(a), we analyzed the testing pixel proportion change along with the training that has both, only one, and none of two segmentor pseudo-labels, i.e., the maximum confidence is larger than as in Eq. (1). We can see that the consensus of the two segmentors keeps increasing, by teaching each other in the co-training scheme for knowledge integration. “Both” low rates, in the beginning, indicate and may provide a different view based on their asymmetric tasks, which can be complementary to each other. The sensitivity studies of using a different number of labeled target domain subjects are shown in Fig. 3(b). Our ACT was able to effectively use . In Fig. 3(c), we show that using more EMD pairs improves the performance consistently.
This work proposed a novel and practical SSDA framework for the segmentation task, which has the great potential to improve a target domain generalization with a manageable labeling effort in clinical practice. To achieve our goal, we resorted to a divide-and-conquer strategy with two asymmetric sub-tasks to balance between the supervisions from source and target domain labeled samples. An EMD scheme is further developed to exploit the pseudo-label smoothly in SSDA. Our experimental results on the cross-modality SSDA task using the BraTS18 database demonstrated that the proposed method surpassed the state-of-the-art UDA and SSDA methods.
This work is supported by NIH R01DC018511, R01DE027989, and P41EB022544.