UPL-SFDA: Uncertainty-aware Pseudo Label Guided Source-Free Domain Adaptation for Medical Image Segmentation

UDA aims to transfer the knowledge learned from labeled source-domain data to an unlabeled target domain. Current UDA methods mainly adapt the model to the target domain in three aspects. The first is image appearance alignment that translates a target-domain image into a source-domain-like image [24, 25, 26, 27], so that the domain gap is alleviated. The second is feature alignment that minimizes the distance of feature distribution between the source and target domains to learn domain-invariant representations [28]. For example, for cardiac image segmentation, Wu et. [29] used Variational Auto-Encoders (VAEs) to align the features in the source and target domains, and Chen et al. [30] used Generative Adversarial networks (GANs) to align the features. The third is output alignment, i.e., using the source model to generate pseudo labels in the target domain for adaptation [6]. However, even relying on unpaired and unsupervised domain translation techniques, these UDA methods require access to source domain images, which is hardly guaranteed at a testing site due to the concerns on privacy, computational cost and bandwidth. Therefore, source-free DA is highly desirable in practice.

Source-Free Domain Adaption (SFDA) deals with domain adaption without access to source-domain data [9, 21, 31, 7]. Yang et al. [31] proposed a Fourier-style mining-guided framework, which comprises a generation stage and an adaptation stage for adapting the source model to the target domain using paired source-like and target images. Sun et al. [9] introduced an auxiliary branch to predict the rotation angle in the target domain. Karani et al. [10] introduced a shallow image normalization network before the segmentation model, and fine-tuned the normalization network in the target domain based on predictions refined by a Denoising Auto-Encoder (DAE). However, these methods require the segmentation model’s structure to be modified in advance to support the auxiliary task and pre-trained with a specified strategy, which is inapplicable to general segmentation models that are unaware of the adaptation process during pre-training. Recently, some methods [32, 12] avoid the coupling between training in the source and target domains, so that the adaptation process does not set a prerequisite for training methods in the source domain, which is more general to arbitrary pre-trained models. Wen et al. [7] proposed a selectively updated Mean Teacher for SFDA, where predictions from a teacher model based on exponential moving average is used to supervise the student. Nado et al. [32] proposed Prediction-Time Batch Normalization (PTBN) that recalculates statistics of batch normalization layers according to the images in the target domain. TENT [12] updates the parameters in batch normalization layers to minimize the entropy of predictions in the target domain. In addition to entropy minimization [12], other loss functions, such as regional nuclear-norm loss with contour regularization [33] and consistency regularization [34], have been proposed for the setting. However, due to the lack of annotations, achieving good performance for SFDA methods is still challenging.

Let $S$ and $T$ be the source and target domains, respectively. Let $\mathbf{X}_{S}$ = $\{(\boldsymbol{x}_{i}^{s},y_{i}^{s}),i=1,...,N_{s}\}$ be the training images and their labels in the source domain, and $\mathbf{X}_{T}$ = $\{(\boldsymbol{x}_{i}^{t},),i=1,...,N_{t}\}$ represent unlabeled images in the target domain for adaptation, where $N_{s}$ and $N_{t}$ are the number of samples in the two domains, respectively. Note that the data distributions in $S$ and $T$ are different, and we assume that the label has the same distribution across the two domains, i.e., the same type of structure for segmentation.

For a general CNN-based segmentation model, it has a feature extractor $g$ and a prediction head $h$, and the parameters of the segmentation model are denoted as $\{\theta_{g},\theta_{h}\}$, where $\theta_{g}$ and $\theta_{h}$ denote the parameters of $g$ and $h$, respectively. As encoder-decoder networks are widely used for medical image segmentation [35, 36], we consider $g$ as an encoder and $h$ as a decoder in this work, respectively. The model is pre-trained in the source domain via:

where $L_{s}$ donates a certain type of supervision loss in the source domain, which might be implemented by fully supervised learning, semi-supervised learning and weakly supervised learning, etc., based on the type of the available labels in the source domain. $\theta_{g}^{0}$ and $\theta_{h}^{0}$ denote the optimized parameter values in the source domain, and they are used as initial parameters for the adaptation process in the target domain.

With the pre-trained feature extractor $g$ and prediction head $h$, the model can be applied to a target-domain image to obtain a prediction as the pseudo label. However, due to the gap between source and target domains, directly applying the pre-trained model will lead to a very low quality of pseudo labels. To improve the quality of pseudo labels for a higher adaptation performance, we propose Target Domain Growing (TDG) to extend the source model, i.e., we duplicate the prediction head (i.e., decoder) $h$ by $K$ times in the target domain, and they are initialized as the pre-trained prediction head with parameter values of $\theta^{0}_{h}$. These prediction heads are connected to the same pre-trained feature extractor $g$ in parallel, as shown in Fig. 1.

Let $h^{k}$ denote the $k$-th prediction head in the target domain. As they have the same initial parameter values with the same architecture, their outputs will be the same for a given input. To obtain diversity, we introduce perturbations for the prediction heads so that they produce different results for more robust ensemble. Specifically, we use random spatial transformation and dropout to improve the diversity of predictions.

First, for an input image $\boldsymbol{x}\in\mathcal{R}^{H\times W}$ in the target domain, where $H$ and $W$ are the height and width, respectively, we send it into the network $K$ times, each time with a random spatial transformation and for a different prediction head $h^{k}$. The segmentation prediction result for the $k$-th head is:

where $\mathcal{T}_{k}$ is a random spatial transformation and $\mathcal{T}^{-1}_{k}$ is the corresponding inverse transformation. $\boldsymbol{p}^{k}\in\mathcal{R}^{C\times H\times W}$ is the output segmentation probability map with $C$ channels obtained by Softmax, where $C$ is the class number for segmentation. In this paper, we set $\mathcal{T}_{k}$ as random flipping, random rotation with $\pi/2$, $\pi$ and $3\pi/2$ for efficient implementation.

Second, we add a dropout layer before each of the prediction head $h^{k}$, so that the prediction heads take different random subsets of the features as input. Due to the image-level and feature-level perturbations, the $K$ predictions are different for an input image. We then average across the $K$ predicted segmentation probability maps for ensemble:

With the average probability prediction $\bar{\boldsymbol{p}}$, we take an argmax to obtain the pseudo label for the input $\boldsymbol{x}$. To reduce noises, we post-process it by only keeping the largest component for each foreground class in segmentation tasks where each foreground class has only one component (e.g, heart structure and fetal brain segmentation in this work). Then the post-processed pseudo label is converted into a one-hot representation, which is denoted as $\tilde{y}\in\{0,1\}^{C\times H\times W}$.

As the domain gap may limit the quality of the pseudo label $\tilde{y}$, directly using $\tilde{y}$ to supervise the network will lead to a limited performance. To deal with this problem, we use the uncertainty information in $\bar{\boldsymbol{p}}$ to identify pixels with reliable pseudo labels and only use the reliable region to supervise the network. To achieve this, we define a binary reliability map $M\in\{0,1\}^{H\times W}$ for $\tilde{y}$, and each element in $M$ is defined as:

where $n$ = 1, 2, …, $HW$ is the pixel index. $c^{*}$ = $\arg\max_{c}(\bar{\boldsymbol{p}}_{c,n})$ is the class with the highest probability for pixel $n$, and $\bar{\boldsymbol{p}}_{c^{*},n}$ represents the confidence for the pseudo label at that pixel. $\tau\in(1/C,1.0)$ is a confidence threshold.

For pseudo label-based self-training, the model may be biased towards its own prediction in each iteration. To avoid this problem, Chen et al. [37] introduced cross supervision where two networks with different predictions guide each other to reduce the bias. However, the use of two networks would increase the computational and memory cost, and it is not suitable for SFDA where only one pre-trained model is provided. Inspired by Chen et al. [37] and to improve the robustness of pseudo label-based SFDA, we introduce Twice Forward pass Supervision (TFS) for robust adaptation.

Specifically, for a batch of data in the training set, before each gradient back-propagation, we perform two consecutive forward passes. We employ the pseudo label $\tilde{y}$ and its associated reliability map $M$ obtained in the first forward pass to supervise the prediction heads in the second forward pass. Let $\boldsymbol{p}^{\prime k}$ denote the output of the $k$-th prediction head in the second forward pass. Due to the use of random spatial transformation and dropout as mentioned above, the outputs of the two forward passes are different despite the same parameter values. Using $\tilde{y}$ to supervise $\boldsymbol{p}^{\prime k}$ can introduce a consistency regularization under perturbations, which improves the robustness of the network. The TFS loss is:

where $\mathcal{L}_{w-dice}$ is the reliability map-weighted Dice loss for a single head. Here we use a Dice-based loss for pseudo label supervision, as Dice loss can better deal with class imbalance in segmentation tasks than cross entropy [38], and the segmentation performance is usually evaluated by Dice.

where $n$ is the pixel index and $\eta=10^{-5}$ is a small number for numeric stability.

Entropy minimization is widely used for regularization in semi-supervised learning [19] and SFDA [39, 13, 40], which improves the model’s confidence by minimizing the entropy of the class distribution in a prediction output. However, existing entropy minimization methods for SFDA are applied to networks with a single prediction head. For our method with multiple prediction heads, enforcing entropy minimization for each prediction head respectively may lead to sub-optimal results when different predication heads obtain opposite results with high confidence. For example, for binary segmentation, when $h^{k}$ and $h^{k+1}$ predict one pixel as being the foreground with probability of 0.0 and 1.0 respectively, both branches have the lowest entropy, but their average has a high entropy. To overcome this problem, we apply entropy minimization to the mean prediction across the $K$ heads:

where $\bar{\boldsymbol{p}}^{\prime}$ is the mean probability prediction obtained by the $K$ heads in the second forward of TFS. Compared with minimizing the entropy of each prediction head respectively, minimizing the entropy of their mean prediction $\bar{\boldsymbol{p}}^{\prime}$ can not only reduce the uncertainty of each head, but also encourage the consistency between them. Thus, it helps to improve the robustness of the network on samples in the target domain.

Our adaptation method adopts a self-training process on unlabeled images in the target domain. Based on the pseudo labels obtained by TDG, the overall loss function for tuning the network with TFS in the target domain is:

where $\lambda$ is a hyper-parameter to control the weight of $\mathcal{L}_{ment}$. Note that there are two forward passes for each parameter update step, where the first forward pass obtains pseudo labels, and the loss function is calculated in the second pass for parameter update with back-propagation.

We used three datasets for experiments: 1) the public Multi-centre, multi-vendor and multi-disease cardiac image segmentation (M&MS) dataset [41], where the images were acquired by devices with four different vendors, 2) an in-house Fetal Brain (FB) segmentation dataset that contains two different MRI sequences, and 3) a public Fetal Tissue Annotation (FeTA) dataset that contains two different super-resolution methods [42]. A summary of these three datasets is listed in Table 1.

To verify the effectiveness of our proposed UPL-SFDA, we compared it with four state-of-the-art SFDA methods: 1) PTBN [32] that updates batch normalization statistics based on unlabeled training images in the target domain without loss functions for optimization; 2) TENT [12] that only updates the parameters of batch normalization layers by minimizing the entropy of model predictions in the target domain; 3) TTT [9] that uses an auxiliary decoder to predict the rotation angle of target-domain images, and the auxiliary task’s loss is used to update the model parameters; and 4) URMA [21] that uses pseudo labels generated by a frozen main decoder to supervise auxiliary decoders. We also compared our method with four naive methods: 1) Source only where the model pre-trained in the source domain is directly used for inference in the target domain, which serves as the lower bound; 2) Target only that uses training images and their labels in the target domain to train a model directly, without pre-training in the source domain; and 3) Fine-tune train and 4) Fine-tune valid that mean the model pre-trained in the source domain is fine-tuned with the annotated training and validation sets in target domain based on fully supervised learning, respectively. In order to investigate the impact of ensembling, we conducted two additional experiments: 1) Source only-Esb that refers to ensemble based on spatial transformations of input images for inference with the pre-trained source model; 2) Ours w/o Esb where our method did not utilize any spatial transformations and made predictions using only one decoder. We implemented all the compared methods with the same backbone, i.e., UNet [35] for M&MS and FB dataset, and 3D UNet [45] for FeTA dataset for a fair comparison.