Uncertainty-Guided Pixel Contrastive Learning for Semi-Supervised Medical Image Segmentation

1 Introduction

Medical image segmentation plays an important role in computer-aided diagnosis system. Supervised learning methods based on deep learning have achieved great performance [Ronneberger et al., 2015; Cao et al., 2021; Li et al., 2021] relying on a large number of labeled data. However, it is difficult to obtain large-scale medical image annotations due to the requirements of professional clinical knowledge and time consumption on data collection and labeling. Semisupervised learning can leverage both labeled data and unlabeled data, which greatly reduces the dependence on annotations. Semi-supervised learning aims to explore the internal information of unlabeled data to improve the performance of the model. Some popular semi-supervised learning strategies include self-training using pseudo-labels [Qiao et al., 2018; Chen et al., 2021a], self-ensembling [Tarvainen and Valpola, 2017; Yu et al., 2019], entropy minimization [Vu et al., 2019], and consistency regularization [Laine and Aila, 2016; Ouali et al., 2020]. The above-mentioned methods utilize unlabeled data during training stage by constructing trusted labels or forcing the prediction consistency when the input data encounters disturbances. However, these methods make the classification of each pixel independent, which ignores the internal correlation between pixels (or features) of images.

∗Corresponding author

In order to strengthen the connection between pixels, some recent works have applied contrastive learning to segmentation tasks [Hu et al., 2021]. The methods of contrastive learning have achieved superior performance in the self-supervised representation learning of natural images. The core idea of contrastive learning is that the representations of similar samples should be alike, and the representations of different kinds of samples should be different. How to define similar samples is the key in contrastive learning. Image-level contrastive learning defines similar samples as different transformations of the same image, and those from different images are defined as dissimilarity. However, similar pixels are densely distributed in the segmentation task. So the definition of dissimilar samples is not suitable for pixel-level contrastive learning. To solve this problem, [Wang et al., 2021] uses segmentation labels to construct contrast samples for supervised segmentation tasks. For unlabeled data, [Chen et al., 2021b] uses the predicted pseudo labels to determine the sample category. [Zhong et al., 2021] uses the spatial consistency of weakly enhanced images to construct similar samples, and construct dissimilar samples by a simple cross-image and pseudo-label weighting heuristic. In fact, using pseudo-labels to construct samples is likely to be inconsistent with the actual semantic categories, which may lead to a noise sampling on contrastive learning. In addition, pixel contrastive learning only establishes the association of local pixels, ignoring the learning of global representation information.

In this paper, our objectives aim to 1) solve the noise sampling problem of contrastive learning using pseudo-labels and 2) strengthen the global representation learning ability of the encoder. To achieve these goals, we propose a contrastive learning method based on uncertainty. Fig. 1 shows the core idea of our method. For unlabeled data, we use the uncertainty map to guide the region of pseudo-labels sampling and reduce the number of wrong samples. Then the sample contrastive loss is calculated to optimize the network and reduce the uncertainty area of predictions. To obtain a better uncertainty map, we design a consistency learning strategy with CNN decoder and Transformer decoder, which can obtain accurate predictions from different views using the structural differences between two decoders. In addition, segmentation models should have the ability to identify geometric transformations. Base on this, we define an equivariant contrastive loss to force the network to learn the identification information of geometric transformations by adding a transformation category prediction in the representation learning stage.

In summary, our contributions mainly include:

• We propose a novel uncertainty guided contrastive learning method, which can effectively alleviate noise sampling from pseudo-labels of unlabeled data.
• A consistency learning strategy for heterogeneous decoders based on CNN and transformer is designed, which can obtain reliable prediction results and uncertainty map by consistency training on unlabeled data.
• We define an equivariant contrastive loss for global representation learning, which equips the model with discrimination ability to distinguish different geometric transformations of images.

2 RelatedWork

2.1 Semi-supervisedMedical Image Segmentation

Without the requirement of large-scale labeled data, semisupervised learning has attracted much attention in medical image segmentation. Existing semi-supervised medical image segmentation methods mainly involves in entropy minimization, pseudo label self-training, collaborative training and consistency learning. Entropy minimization [Vu et al., 2019] suggests that high-quality prediction results should have a low entropy, and hence it conducts model learning by minimizing the information entropy of the prediction probability distribution. Pseudo label self-training [Chen et al., 2021a] performs class supervised learning by predicting pseudo labels for the unlabeled data. Co-training [Qiao et al., 2018] assumes that there are multiple decision views containing complementary information, and designs different classifiers to learn different views to promote segmentation performance. Consistency learning [Verma et al., 2019; Laine and Aila, 2016; Tarvainen and Valpola, 2017; Ouali et al., 2020] makes an assumption that even if an image sample encounters some disturbances, say input disturbance or model disturbance, the prediction results from the sample should not change. Motivated by such intuition, these methods conduct model training by encouraging consistent prediction of unlabeled disturbed samples. Inspired by collaborative training and consistency learning, we propose to characterize the complementary information of data from different views, using the structural differences between CNN and transformer, and apply consistency constraints to train the model.

2.2 Contrastive Learning

In image-level representation learning, contrastive learning can make full use of unlabeled data to learn effective visual representation, in which the core idea is to strengthen the discrimination of the learned visual representation by narrowing similar pairs (positive) and separating dissimilar pairs (negative) based on some similarity constraints. The key point of image-level contrastive learning is how to construct contrastive samples. A feasible solution is proposed in [He et al., 2020], which increases the number of contrastive samples by introducing memory bank and momentum contrast.

Recently, some works [Chaitanya et al., 2020; Wang et al., 2021; Zhong et al., 2021; Hu et al., 2021] have been proposed to extend contrastive learning from image-level to pixel-level for image segmentation. The main idea of pixel-level contrastive learning is to construct pixel sample pairs with the help of segmentation labels. For unlabeled data, sample pairs are constructed by using pseudo labels or spatial structure. Nevertheless, these methods may encounter the problem of noise sampling during the process of constructing sample pairs. To alleviate this problem, we suggest using prediction uncertainty to guide sample sampling and reduce the number of noise samples. In addition, pixel-level contrastive learning lacks the capture ability of global representation, which urges us to impose a constraint of prior knowledge in representation learning for segmentation task.

2.3 Uncertainty Estimation

In semi-supervised learning, uncertainty can be used to evaluate the quality of model predictions for better use of unlabeled data. The measure methods of estimating uncertainty mainly include 1) using information entropy of the prediction probability distribution, 2)using the deviation of multiple predition results of the same input under different disturbances [Yu et al., 2019], and 3) calculating the variance of different prediction results of the same input [Zheng and Yang, 2021]. However, these methods are time-consuming and lack of reliability. In our method, we estimate the uncertainty by calculating the entropy of the average probability distribution obtained by different predictors to overcome these problems.

3 Methodology

Given a label dataset DL = {(xi, yi), i = 1, . . . ,N} and an unlabeled dataset DU = {xj, j = 1, . . . ,M}, where M ≫ N, the images in DL and DU firstly go through a geometric transformation and then feed into an encoder network to extract multi-scale features. After that, these features will be sent to the following three branches, including consistency learning branch, uncertainty-guided contrastive learning branch, and equivariant contrastive learning branch. For the consistency learning branch, we propose a heterogeneous consistency network to predict the segmentation results, which is driven by a supervision loss ℓsup and a consistency loss ℓcon. ℓsup is calculated by ground truth of DL, and ℓcon is calculated by prediction consistency of DU. For the uncertainty-guided constrastive learning branch, we build and maintain a memory queue to preserve enough samples for constrastive learning. The selection of samples in the memory queue depends on the labels of DL and certainty labels of DU. For the selected samples, we impose a pixel-level contrastive loss ℓc to make the pixels of the same class close to each other and the pixels of different classes far away from each other. For the equivariant contrastive learning branch, we perform geometric transformation category prediction on all labeled and unlabeled data, and design an equivariant contrastive loss ℓe to force the encoder to be robust to geometric transformation. For ease of understanding, Fig.2 gives an illustration of the overall architecture and the training process of our proposed method. In summary, the total objective of

In this paper, we set λ1 = λ2 = 0.1 and λt is a temperature parameter that increases from 0 to 0.01. The following shows the above-mentioned three branches in details.

3.1 Consistency Learning Between Decoders

In the consistency learning branch, we design a simple yet effective network structure to achieve the following two goals:

1) using unlabeled data to promote the learning of segmentation network, and 2) obtaining reliable uncertainty estimation from network outputs. It is demonstrated that using cotraining strategy can obtain better segmentation performance, of which the core idea is to make different classification predictions from different views and then regard the differences of predictions as the measurement criterion of uncertainty estimation. Instead of using the same architecture that requires adding some disturbances for co-training, inspired by [Luo et al., 2021], we adopt a simple yet effective scheme that takes advantage of the congenital differences between transformer decoder and CNN decoder. Specifically, we construct a heterogeneous predictor to constrain the two decoders to generate consistent predictions. And then the entropy of the mean prediction is used to estimate the uncertainty map.

Patch Embedding and Position Encoding. We choose two different decoders f t θ (·) and f c θ (·) from Swin-UNet and UNet. Through the encoder, we can obtain a set of features{fi, i = 0, . . . , 3}. Before inputing into f t θ (·), we need to reshape the feature fi ∈ RHi×Wi×Ci into a sequence of flattened patches xip ∈ RP 2 i ×Ci and embed them into a Ddimensional space using a learnable linear projection E. To preserve the spatial information, we add the absolute location encoding to the embedded patches:

where E ∈ R(P 2 i ×Ci)×D is the patch embedding projection, and Epos ∈ RN×D denotes the absolute location encoding.

Consistency Learning. Given an input image, we can get two predicted probability distributions pt and pc from two decoders. For labeled data, we use ground thuth to calculate the supervision segmentation loss:

where LCE, LDice are the cross-entropy loss and dice loss, and y is ground truth of labeled data. We use f c θ (·) as the main predictor, so α is set to 0.4. In the inference stage, prediction from CNN branch is the final result. For unlabeled data, consistency loss is calculated as follows:

where Ldis is a distance measure between two output probability distributions. In this work, we choose to use mean squared error (MSE) as a distance measure.

3.2 Uncertainty-Guided Contrastive Learning

Image segmentation task generally involves with cross entropy loss for conducting pixel-level classification. However, such loss makes the classification of each pixel independent and hence ignores the relationship between the pixels. To solve this problem, we design a pixel-level contrastive learning mechanism that classifies the pixels of the same category (semantic label) as positive samples, and the pixels of different categories as negative samples. The relationship between pixels is established by reducing the distance between positive samples and enlarging the distance between negative samples in an embedding space. In order to effectively use unlabeled data for pixel-level contrastive learning, we estimate the uncertainty of unlabeled pixels and the pixels with a higher certainty are selected as anchor points for contrastive learning.

Mask Uncertainty Region. We choose the predictive entropy as the metric to approximate the uncertainty. Specifically, we first calculate the average probability distribution of the prediction results ˆ p = (pc + pt)/2, and then calculate the entropy for the probability distribution of each pixel in the channel dimension. It can be summarized as:

where ϵ is a very small constant to avoid singularity. We believe that the prediction with large entropy is uncertain in category. When calculating the pseudo labels, those uncertainty predictions are removed as a non-sampling region, and then the determined pseudo labels are obtained:

where H is a threshold to mask the uncertain labels, and yp is the final certainty pseudo label.

Anchor Sampling. We use the labels of labeled images and the certainty pseudo labels of unlabeled images as the basis for the use of contrastive samples. Because the original image resolution is too large, the cost of contrastive learning in the original image size is expensive and the prototype vectors of pixels contain less semantic information. Therefore, we use contrastive learning in the feature space with low resolution. Firstly, the features extracted from the encoder will be embedded into the D-dimensional space, where each Ddimensional feature vector represents the prototype vector of the pixel. Then, the labels are downsampled to the same resolution, the category is specified for each prototype vector, and the vector in the uncertainty region is not sampled. We adopt the strategy of random sampling with a fixed number of samples for each category. If the number of samples from same category is small, we will sample anchors from other categories. The number of contrastive negative samples greatly affects the performance of contrastive learning, but a large number of negative samples will produce a lot of overhead. A better solution is to use a fixed size external storage to store the sampled samples and update the storage content with the training. In our method, we set up a memory queue to store the collected samples. In each iteration, the randomly selected samples are used as anchors to calculate the contrastive loss, and then they are updated to the memory queue.

Pixel Contrastive Loss. The prototype vectors and their category of pixels are saved in the sample queue. We use the popular InfoNCE [van den Oord et al., 2018] loss function to calculate the contrastive loss. In each iteration, we randomly sample M anchors and calculate the contrastive loss for each anchor. Then average the loss of all anchors as the overall contrastive loss. The specific calculation is as follows:

where Pi and Ni denote prototype vector's collections of the positive and negative samples for pixel i. vi is the prototype vector of pixel i, v+ i is a positive prototype vector, vi is an negative vector and τ is a temperature hyper-parameter.

3.3 Equivariant Contrastive Loss

For conducting contrastive learning, some previous work constructs positive samples by different transformations of the same image. However, some transformations do not accord with the prior knowledge of segmentation tasks[Dangovski et al., 2021], such as geometric transformation. In this paper, we suggest that the effective feature representations required by segmentation task should be equivariant (or discriminative) to different geometric transformations.

Based on the above, we consider adding the equivariant contrastive loss to the representation learning of the segmentation model to learn the global information. Specifically, we define the segmentation model as encoder-decoder form: f(xi) = fγ(fθ(xi)). For an image xi, when it passes through some geometric transformation G(·), the corresponding segmentation result will also change, that is:

Then, we can deduce:

Therefore, we can explicitly strengthen the learning of this geometric transformation information in fθ(·). We add a classification predictor pϕ(·) to predict the discrimination results of geometric transformation. Our equivariant contrastive loss function is as follows:

where the geometric transformation Gi (·) represents the fourfold rotation in this paper, so C = 4.

4 Experiments

4.1 Experimental Setup

Datasets and Metircs. We validate the proposed method on two public datasets:

• ACDC dataset [Bernard et al., 2018] contains 200 annotated short-axis cardiac MR-cine images from 100 patients. We divide the dataset in a ratio of 7:3 to obtain the training set and verification set. According to different semi-supervised experiment settings, 136 images from 7 patients and 68 images from 3 patients in the training set are labeled respectively. See SSL4MIS 1 for details.
• ISIC dataset [Codella et al., 2018] includes 2594 dermoscopy images, and we use 1815 images for training and 779 images for validation. In training set, 5% (91) and 10% (181) images are labeled for different semisupervised experiment settings.

All images in both datasets are resized to 224 × 224 to meet the input requirements of the proposed method. We use standard data augmentation to enlarge training set, including random cropping, random rotating, random flipping and color jittering. In our method, the category of transformation will be recorded to calculate the ℓe. In order to evaluate the performance of our method, we select Dice Coefficient (denoted as Dice) and Jaccard Index (denoted as Jaccard) as evaluation metrics.

1 https://github.com/HiLab-git/SSL4MIS

Implementation Details. For fair comparisons, all the methods used in the experiments choose UNet as the benchmark architecture for image segmentation. We use ResNet-50 to replace the encoder part of UNet and initialize its parameters with the weights pre-trained on ImageNet. We adopt SGD as an optimizer with a weight decay of 0.0005 and a momentum of 0.9. The initial learning rate is set to 0.01, which will reduce to 0.001 by polynomial scheduler strategy during training. We implement the methods using PyTorch library and train them on a NVIDIA RTX 2080Ti GPU. The batch size is set to 16, where 8 images are labeled. All methods perform 6000 iterations during training.

4.2 Quantitative Comparison

Compared Methods. We compare our method with some recent semi-supervised segmentation methods including: Meat-Teacher (MT) [Tarvainen and Valpola, 2017], Entropy Minimization (EM) [Vu et al., 2019], Uncertainty-Aware Mean Teacher (UA-MT) [Yu et al., 2019], Deep Co-training (DCT) [Qiao et al., 2018], Cross-Consistency Training (CCT) [Ouali et al., 2020] and Cross Pseudo Supervision (CPS) [Chen et al., 2021a]. For all comparison methods, we adopt the official hyper-parameter settings.

Main Results. Table 1 shows our quantitative comparative experimental results on ACDC and ISIC datasets. The first row represents the performance of the baseline model trained with labeled data only. Compared with the baseline model, our method can effectively use unlabeled data to achieve great performance. In different datasets and different semisupervised settings, our proposed method obviously outperforms the comparison methods. Especially, when only 68 labeled images are used in ACDC dataset, our method improves Dice by more than 3% compared with the other methods.

Visual Comparisons. Fig.3 shows some visual comparisons between different methods when using 136 labeled images on ACDC dataset. We chose two methods MT and CPS which performed better in the experiment for comparison. Compared with MT and CPS, our method has better prediction results and less false predictions.

4.3 Ablation Study

Table 2 shows the results of ablation experiments of our method on ACDC dataset with 136 labeled images. We choose the UNet model that only uses labeled data for supervised training as the baseline (first row), and gradually increase the proposed components to prove their effectiveness. Besides, we further add two additional comparison settings, including 1) using only labeled data for contrastive learning (ℓsup c ) and 2) using pseudo labels for contrastive learning (ℓpseudo c ) to demonstrate the effectiveness of our proposed uncertainty-guided contrastive learning method. The experimental results show that each part of our proposed method has a positive impact. The introduction of contrastive learning effectively establishes the relationship between pixels and improves the performance of the model. Compared with the pseudo label scheme, the proposed method makes full use of unlabeled data and hence brings significant performance improvements (Dice increases by about 1%).

Contrast on Different Feature Scales. The resolution of different features has an important impact on the selection of contrastive learning samples. To find an appropriate feature scale, we explore the effect of different scales for contrastive learning on the ACDC dataset. Table 3 shows the results of contrastive learning under four different feature scales. As we can see, the performance of contrastive learning under low resolution (Conv4) is poor, which might be caused by the semantic inconsistency from label downsampling. The higher resolution feature (Conv1) also brings performance degradation. The potential reason is that the high-resolution pixel vectors contain less semantic information. Therefore, we believe that using middle-level features for contrastive learning can bring better segmentation performance.

Visualization of Features. In Fig.4, we use t-SNE algorithm to reduce the dimension of pixel features for visualization. From left to right, they are the results of training without contrastive learning, contrastive learning using pseudo labels and the proposed method. Compared with the first one, the proposed method can equip the pixel representations with better intra-class compactness and inter-class separability, which indicates the effectiveness of contrastive learning for segmentation task. Compared with the second one, our method has better aggregation results, the potential reason of which is that our method can reduce the possibility of noise sampling.

5 Conclusion

We propose an uncertainty-guided pixel contrastive learning method for semi-supervised medical image segmentation, which uses uncertainty to solve the noise sampling problem of unlabeled data in pixel contrastive learning. To estimate uncertainty, a heterogeneous consistency learning strategy is elaborately designed based on the decoders of CNN and Transformer. In addition, we construct an equivariant contrastive loss to strengthen the global representation learning ability of our model. Extensive experiments demonstrate that our method can achieve the state-of-the-art performance.

Acknowledgements

This work was supported in part by the Natural Science Foundation of China under Grant 61976145 and Grant 61802267, in part by the Shenzhen Municipal Science and Technology Innovation Council under Grants JCYJ20210324094413037 and JCYJ20190813100801664, and in part by the Guangdong Basic and Applied Basic Research Foundation 2021A1515011318.