{'fraction_non_alphanumeric': 0.031818181818181815, 'fraction_numerical': 0.011499321573948439, 'mean_word_length': 4.806775654914319, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Uncertainty-Aware Deep Co-training for Semi-supervised Medical Image Segmentation

Semi-supervised learning has made significant strides in the medical domain since it alleviates the heavy burden of collecting abundant pixel-wise annotated data for semantic segmentation tasks. Existing semi-supervised approaches enhance the ability to extract features from unlabeled data with prior knowledge obtained from limited labeled data. However, due to the scarcity of labeled data, the features extracted by the models are limited in supervised learning, and the quality of predictions for unlabeled data also cannot be guaranteed. Both will impede consistency training. To this end, we proposed a novel uncertainty-aware scheme to make models learn regions purposefully. Specifically, we employ Monte Carlo Sampling as an estimation method to attain an uncertainty map, which can serve as a weight for losses to force the models to focus on the valuable region according to the characteristics of supervised learning and unsupervised learning. Simultaneously, in the backward process, we joint unsupervised and supervised losses to accelerate the convergence of the network via enhancing the gradient flow between different tasks. Quantitatively, we conduct extensive experiments on three challenging medical datasets. Experimental results show desirable improvements to state-of-the-art counterparts.


Introduction
Semantic segmentation is a fundamental computer vision task that is particularly important for medical image analysis [1]. It serves as a preprocessing step for the treatment of different medical conditions. In recent years, convolutional neural network (CNN) based methods have achieved tremendous improvement [2], [3]. With abundant pixel-wise labeled data, CNN can easily achieve outstanding performance relying on great nonlinear fitting capability [4]. However, pixel-wise annotated images are expensive to obtain especially in the medical domain. It is time-consuming and extremely relies on experienced experts. Fortunately, semi-supervised methods leverage not only labeled data but also unlabeled data [5], freeing the researchers from labeling work. Therefore, semi-supervised learning has attracted great attention [6], [7], [8].
The current dominant semi-supervised learning methods in deep learning are self-training [9], consistency regulation [10], pseudo labeling [11], adversarial learning [12] and so on. All the above approaches are based on the assumption that semi-supervised methods rely on intrinsic properties of the dataset distribution rather than individual images. So, parameter optimization can process with a combination of annotated images and unannotated images rather than labeled images only. For example, self-training [10] selects the most reliable prediction of the current model as the label (pseudo label), for the complement of the labeled dataset. Other approaches attempt to exchange internal information among an ensemble of models, aimed at increasing their consistency, such as co-training [13]. These approaches are all constructive to a certain degree. But how to leverage unannotated data more effectively is still one of the most concerning issues in these semi-supervised learning methods.
Most of the existing semi-supervised methods leverage the prediction of unannotated data, but the quality Figure 1: Overview of the uncertainty-aware co-training algorithm. On the basis of training two deep CNN models simultaneously with two different sets of the labeled datasets and the same unlabeled dataset, we bring in Monte Carlo Sampling to acquire uncertainty, which is acting on the loss function in both the supervised learning stage and the unsupervised learning stage. of the prediction is not guaranteed. Optimizing model parameters via unreliable predictions in unsupervised learning is not convinced, even towards wrong results. Here, the major challenge is to assess the reliability of the prediction. To address this issue, we propose an uncertainty-aware semi-supervised semantic segmentation framework, as illustrated in Fig. 1. Uncertainty is an estimation of prediction confidence [14], which is leveraged as prior information to guide network learning by adjusting loss weight dynamically in two separate strategies in our framework. It forces the models to focus on the valuable region according to the characteristics of supervised learning and unsupervised learning. On the one hand, it guides networks to pay more attention to the pixels with high prediction reliability (low uncertainty) in unsupervised learning; on the other hand, it guides networks to concentrate more on ambiguous pixels (high uncertainty) in supervised learning. In addition, owing to its success in semi-supervised learning, deep co-training is utilized in our proposed framework. The general principle of co-training is to independently train models with the labeled data and enforce their predictions to agree on the labeling of the unlabeled data. We separate the labeled dataset into complementary subsets and use adversarial examples [15] to enforce the diversity. Moreover, to promote the predictions consistency among different models, we add ensemble agreement loss to guide models training. We joint unsupervised and supervised losses to accelerate the convergence of the network via enhancing the gradient flow between different tasks.
In this paper, we explore the uncertainty information to enforce the performance of networks both in supervised learning and unsupervised learning. With the guidance of the estimated uncertainty in both supervised and unsupervised stages, our uncertainty-aware deep co-training method achieves the state-of-the-art performance in the semi-supervised tasks on three public medical datasets. Moreover, a comprehensive set of experiments are conducted which demonstrate the potential of our approach for segmenting different types of medical images. Our experiments analyze the impact of various elements of the method.

Semi-supervised segmentation
Semi-supervised learning is an indispensable part of deep learning theory [16]. Recently, many works have proposed using semi-supervised approaches in medical image segmentation to segment human organs [17]. An innovative semi-supervised segmentation approach is presented by Mohamed et al. for efficient segmentation of lung CT scans from Convid-19 patients. Yu et al. [18] proposed a teacher-student model to segment 3D left atrium via self-ensembling. Li et al. [10] proposed a semi-supervised segmentation method for skin lesions by enforcing the consistency between the student and teacher model. Fang et al. [19] used a co-training framework to boost each sub-model for kidney tumor segmentation. And they applied adversarial training [20] to their network to constrain the models to output invariant results over different perturbations on input data.
Similar to co-training, Peng et al. [13] presented a method based on an ensemble of deep segmentation networks. They trained different models with corresponding subsets of the annotated data and leveraged non-annotated images to exchange internal information among sub-models. The difference between these semisupervised methods is how they utilize unlabeled data and the way they relate to supervised algorithms [5]. Most of the approaches utilize predictions of unlabeled data for consistency training [21]. However, the model's prediction of unlabeled data used to exchange information is unstable. We need to add components to enable us to learn the specific confidence situation of each pixel of the prediction [22]. So we introduce uncertainty estimation [23] in our approach.

Uncertainty Estimation
Uncertainty estimation plays a pivotal role in reducing the impact of stochastic during both optimization and decision-making processes [24]. Knowing the confidence (uncertainty) with which we can trust the neural networks' predictions is essential for decision making [25]. A mathematically grounded framework is offered by Bayesian probability theory to learn about the uncertainty of a data-driven model, called Bayesian Neural Network(BNN). BNN attempted to learn a distribution over each of the network's weight parameters [14]. However, Bayesian inference is intractable on computation in practice.
Deep ensemble [26] is an alternative to BNN, which observes the variance of many trained models to estimate uncertainty. In addition, several works applied different uncertainty estimation methods to the task of semantic segmentation, including Stochastic Batch Normalization [27], Multiplicative Normalizing Flows [28] and so on. Upadhyay et al. [29] proposed a GAN-based framework that estimates the per-voxel uncertainty in the predictions on magnetic resonance imaging (MRI). Alex [14] studied models under the framework with per-pixel semantic segmentation and proposed new loss functions based on uncertainty, interpreted as learned attenuation. Wu et al. [21] claimed that deep models require extra components to obtain performance gains. And Yu et al. [18] introduced uncertainty to ensemble models by using Monte-Carlo Sampling. It has been proved that the Monte Carlo Sampling has much superiority over conventional methods in the estimation of uncertainty, especially of complex measurement systems' output [25].

Problem formulation
Since manual annotation is often time-consuming and expensive, only a small fraction of images in datasets can have full pixel-wise labels. So our proposed method aims to learn more information from both labeled and abundant unlabeled images by using the uncertaintyaware deep co-training methodology. We formalize the problem of medical image semantic segmentation as follows.
Given labeled dataset D l = {(x l 1 , y l 1 ), ..., (x l m , y l m )}, which contains m labeled examples, each example comprised of an input image x l i with dimensions H × W and its corresponding pixel-level label y l i ∈ R H×W×C , where C is the number of classes and H × W is spatial dimension . In semi-supervised setting, we also have unlabeled dataset D u = {x u 1 , ..., x u n } including n unlabeled images, with n m. The purpose of semi-supervised segmentation is to train a segmentation model f of parameter θ with D = D l ∪ D u , which map each pixel of an input image to its correct label.

Proposed approach
We train two models f (θ 1 ), f (θ 2 ) jointly with labeled dataset D l and unlabeled dataset D u in a collaborative manner. Motivated by performance improvement from the uncertainty estimation in the Bayesian networks, we employ Monte Carlo Sampling [14] to estimate uncertainty. Specifically, we perform T stochastic forward passes on two models under random dropout. According to this, we obtain a set of predictions and choose the predictive entropy metric to approximate the uncertainty. The predictive entropy can be summarized as: where P c t is the probability of the c-th class in the tth time prediction. And the uncertainty is estimated in pixel level, the uncertainty of the whole volume U is {u} ∈ R H×W×D .
Following co-training methods for semi-supervised semantic segmentation, we employ a loss function composed of a weighted sum of three separate terms to train ensemble models.
(2) Uncertainty is used to impact the weight of the first two parts of the loss function. The details are explained in the following subsections.

Supervised learning
As in standard multi-view learning [30] approach, we separate the labeled dataset D l to complementary subsets D l 1 and D l 2 . Training two models individually with corresponding labeled subsets makes each model have a certain fit to the distribution of the data. And also, we can ensure the diversity of the two models.
We employ E-Net [31] as our network backbone. To adapt the E-Net as a Bayesian network to estimate the uncertainty, one dropout layer is added between the encoder and decoder. In the network training and uncertainty estimation, we turn on the dropout layer. Since we don't need to estimate uncertainty in the testing phase, we turn off the dropout layer.
Through the forward propagation of each model T times, we can get T different results as shown in Fig. 2. Therefore, we obtain a set of Softmax probability vectors for each pixel in the input. By using predictive entropy as the metric, we can approximate uncertainty maps. We express the uncertainty map as the model confidence of each pixel prediction. If a pixel has a large fluctuation during the prediction process of T forward propagation, the uncertainty of it will be large.
With the guidance of the estimated uncertainty U sup 1 and U sup 2 , we find out the unreliable prediction in images and make models learn more from these pixels from corresponding labels. We design the uncertainty-aware supervised loss L sup as the pixel-level crossentropy loss of the two models:

Unsupervised learning
Semi-supervised learning means not only exploiting labeled images, unlabeled images D u are more important for the learning process. The strategy of unsupervised learning wants the segmentation networks to output similar predictions for the same unlabeled image with perturbations. The deep co-training method minimizes the distance between the class distributions predicted by different models. However, the prediction results obtained through one forward pass are contingent, which will limit models' ability to fit the unlabeled data distribution.
So, we also introduce the concept of uncertainty in unsupervised learning. Same as the fully supervised phase, we perform T stochastic forward passes on each model under random dropout as shown in Fig. 3. After leveraging predictive entropy as a metric, we get two uncertainty maps approximated by models. Different from the supervised stage, we combine the results of the two models here. We calculate the average of the two uncertainties and take it as a weight coefficient and multiply it on the loss function. Since the calculated uncertainty fluctuates greatly, we normalize it in this way: where β and C are constants used for normalization. In order t to make the uncertainty change within the controllable range, we set β = 0.7, C = 2.
In supervised learning, the models should pay more attention to the pixels with higher uncertainty. But here in unsupervised learning, models should learn more from lower uncertainty pixels as there is no label in this stage. The estimated uncertainty guides models concentrating more on reliable (low uncertainty) predictions. We add uncertainty U un 1 , U un 2 estimated by two models together, and normalize it as U un . The uncertainty-aware ensemble agreement loss we leverage here is Jensen-Shannon divergence: Figure 3: Illustration of unsupervised learning. Two networks output predictions based on the knowledge from each view. We leverage the ensemble agreement loss L con to encourage the predictions of models to consistency with the guidance of the estimated uncertainty. The final uncertainty is the sum of the uncertainties generated by two models using Monte Carlo Dropout. Each model performs T stochastic forward passes to get an uncertainty map.

When to introduce uncertainty?
Introducing uncertainty into models can improve the performance of the network, but introducing it at unappropriate time will crush the entire framework.
In the supervised stage, theoretically, the models don't need prior knowledge to avoid turning in the wrong direction. So the sooner uncertainty is introduced to guide network learning, the better performance we can get. We verify this by doing experiments of taking uncertainty in concern at different epoch. The results are shown in Fig. 4. Considering the performance in both DSC and HD, we leverage uncertainty from the beginning in the supervised stage.
In the unsupervised stage, the models' prediction error for unlabeled images is large at the beginning of experiments. This will lead to a large fluctuation range of uncertainty, which is not conducive to allowing the model to focus on where it should learn more. So, we tested adding uncertainty to the unsupervised stage at different epoch, and the final results are shown in Fig. 4. And finally, we learned that introducing uncertainty in the unsupervised stage requires that models have a certain understanding of the target distribution.
Synthesize all factors, we introduce uncertainty to the supervised stage at epoch 0, and epoch 20 for the unsupervised stage in all experiments of our approach.

Adversarial learning
Having diversity between models in the ensemble is indispensable. In co-training, diversity is essential for avoiding the collapse of decision boundary and the models can learn from each other during training. The deep co-training method uses the approach proposed by Qiao et al. for image classification and augments the dataset with adversarial examples. The adversarial examples are generated from both labeled and unlabeled data.
One model uses adversarial examples teaching other models in the ensemble. In the original deep co-training method, the diversity loss is: where H(.) refers to cross entropy and x is an input image. The g i (x) is an adversarial example target on model f i (.; θ i ). So g 1 (x) is an adversarial example for model 1, and we can easily get that f 1 (x; θ 1 ) f 1 (g 1 (x); θ 1 ). In the process of minimizing the first term of Eq. 6., it will make f 1 (x; θ 1 ) approximately equal to f 2 (g 1 (x); θ 2 ). Combining the above relations, we can obtain f 1 (g 1 (x); θ 1 ) f 2 (g 1 (x); θ 2 ). Applying the same idea to model 2, we can conclude that the models have divergence in predicting adversarial examples. Adversarial examples can be obtained by adding small perturbations to input images. As described in [13], we follow use different schemes to generate these examples based on the source of the image x. We apply the Virtual Adversarial Training (VAT) method when x is an unlabeled image from D u ; and we use the Fast Gradient Sign Method (FGSM) when x in labeled datasets.

Dataset and metrics
We evaluated our method on three medical image segmentation benchmark publicly available datasets: Automated Cardiac Challenge (ACDC) [32], Spinal Cord Gray Matter Challenge (SCGM) [33], and Spleen sub-task dataset of the Medical Segmentation Decathlon Challenge [34]. subjects (150 scans) for training and 25 subjects (50 scans) for testing. All short-axis slices within 3D-MRI scans were resized to 256 × 256 as 2D images.
SCGM dataset. : The Spinal Cord Gray Matter Challenge was organized to test the different capabilities of various methods. The dataset in this challenge is a public-available collection of multi-vendors, multicenter MRI. All data was acquired at 4 different sites: University College London, Polytechnique Montreal, University of Zurich, and Vanderbilt University. It contains 80 healthy subjects (age range of 28.3 to 44.3 years) with 20 subjects from each center. The training set contains 40 labeled scans, each annotated slice-wise by 4 independent experts. The ground truth mask is obtained by majority voting. In our work, we leverage 30 images from the first center as a labeled dataset, and 465 images from all centers as an unlabeled dataset. The test set contains 264 labeled images from center 3 and center 4. The slices are center-cropped to 200 × 200 pixels.
Spleen dataset. : Spleen dataset is one of the ten sub-tasks of the Medical Segmentation Decathlon Challenge [34]. It consists of patients undergoing chemotherapy treatment for liver metastases as a pub-licly available dataset. The dataset includes 61 portal venous phases CT scans (only 41 were given with ground truth). The ground truth segmentation was generated by a semi-automatic segmentation software and identified by an expert abdominal radiologist. Each slice obtained from the CT scans is resized to 256 × 256 pixels. We split the dataset into labeled (4 patients), unlabeled (32 patients), and validation image subsets( 5 patients).
Dice similiarity coefficient(DSC). : DSC measures the overlap between the predictions of the model S and the ground truth G: Hausdorff distance(HD). : HD is a boundary distance metric which measures the largest distance (in mm) between a point in prediction S and the closest point in the ground truth G:

Implementation Details
For the sake of speed and accuracy, we leverage the well-known lightweight E-net [31] as our basic segmentation network. This architecture is one of the most popular models for segmentation. To adapt the E-Net as a Bayesian network to estimate the uncertainty, one dropout layer with a dropout rate of 0.5 is added between the encoder and decoder of the E-Net. All the datasets are applied random rotation, random crop, and flip as augmentation strategies. We set to 0.03 in FSGM and 10 in VAT. The learning rate decreased with a factor of 10 every 90 epochs, and it is 10 every 100 epochs in SCGM experiments. To describe the level of supervision, we vary the ratio l a (0 ≤ l a ≤ 1) of labeled samples in our experiments.
For all experiments, the setting of λ is with a dynamic strategy. We used a ramp-up strategy followed the Gaussian ramp-up curve. And for all the hyperparameter λ settings, we follow [13]. The framework is implemented with Pytorch library, trained on one NVIDIA 2080Ti GPU.

Comparison with semi-supervised methods
Our uncertainty-aware co-training method is compared against several recent state-of-art semi-supervised methods in the medical domain: Mean Teacher. : Mean Teacher (MT) is a method using multiple deep CNNs for semi-supervised segmentation. And it is an effective approach that averages model weights instead of predictions.
Deep-Co-trianing. : Deep-Co-training (CT) shares information between simultaneously trained models, while preserving their diversity.
We also compare our method with the full supervised method of the same amount of labeled images as the semi-supervised method, referred to as "Part". "IND" means the performance of individually-trained models (Independent). We report the performance of a fully supervised baseline ("Full") which means training a single model with all available datasets.
For these baselines, we follow the same learning rate decay, weight scheduler, data augmentation setting, and optimization as for our method. For MT and UAMT, we set EMA (Exponential Moving Average) parameters α to 0.99. We report both the average performance of individual models ("avg") and the performance of combining the predictions of all models using a voting strategy    ("vot") called ensemble soft voting. This strategy leads to a higher accuracy than the prediction of individual models for both DSC and HD. In order to avoid contingency, we compute the average and standard deviation over three runs with different random seeds.
All the compared baseline results reported in our paper are reimplemented by us.

ACDC dataset
Our uncertainty-aware method is first evaluated on the ACDC dataset with l a = 0.2. Tab. 2. presents the quantitative results. Compared with other methods, our approach achieves higher overall performance on the test dataset in terms of DSC and HD. Especially, our approach achieves improvement over co-training with an overall performance increase of 2.29% in DSC and 0.49mm in HD, and the standard deviation of multiple trials (three runs) of the results of our approach is the smallest. With only 20% labeled images, we are only 1.8% in DSC and 0.39mm in HD behind the 100% label ratio supervision. Some examples of results from the test dataset are shown in Fig. 5. We can see our method gives contours closer to ground truth (GT), with more accurate segmentation in details between different regions.
Also, we evaluate how label ratios impact results in a dual-view setting. Tab. 3. shows the results for different labeled data ratios: 5%, 10%, 20%, 30%, 40%, 50%. We can easily find that as the label ratio increase, mean DSC values increase sharply, and mean HD values decrease. In all cases, our approach leads to a better performance in DSC and HD than training models separately and deep co-training method. Our approach has more obvious advantages with a low label ratio. With 5% label ratio, our method outperforms Deep co-training 2.73% in DSC and 1.34mm in HD.

SCGM dataset
We evaluate our method on the task of segmenting spinal cord grey matter in images from the SCGM dataset. The SCGM dataset is from four different clinical centers, so different parameters are applied in collecting the MRI images. We only used a few labeled images (i.e., only 30 images from one center), and test sets are from the other centers. The labeling ratio is about 6.5%. Due to the different data sources, it is more difficult to extract the image features of the samples. This also leads to low segmentation accuracy of semisupervised methods in this task. The use of uncertainty solves this problem well, models can learn effective semantic features from different centers according to the uncertainty. The results are summarized in Tab. 4. Our approach gives a mean DSC of 5.96% and HD of 6.35mm better than the best baseline (deep co-training). And our method gives a mean DSC of 12.94% and HD of 23.26mm better than the performance of the fully supervised baseline. Fig. 6. shows the segmentation results on the test dataset. In some difficult-to-recognize images, fully supervised training cannot even segment the lesion area, but our approach can complete the segmentation task to a certain extent.

Spleen dataset
We also validate the effectiveness of our uncertaintyaware deep co-training method on the task of segmenting spinal cord grey matter in images from the Spleen dataset. We repeated our experiments on the Spleen dataset consisting of 2D slices of CT scans resized to a resolution of 256 × 256 pixels. Tab. 5. summarizes the experimental results. We see that the Deep co-training method's performance is almost the same as the fully supervised baseline. In case of better stability (standard deviation) of the results, the accuracy has been improved and surpassed full supervision after taking uncertainty into concern. Our method improves accuracy over DCT-seg from 94.09% to 95.45% and HD from 6.86mm to 5.57mm. Examples of segmentation results obtained by tested methods are given in Fig. 7.
Semi-supervised learning pursues using less labeled images to achieve better segmentation accuracy. We show the performance of our approach, DCT-seg, and individually trained models(independent) on smaller label ratios: 5% and 10%. Tab.6. gives the mean DSC with standard deviation. Our approach is still in the leading position among all semi-supervised methods.

Abalation analysis of our method
We do ablation studies to prove that our method works in both supervised and unsupervised stages on the ACDC dataset. The results are shown on Tab.6.
In supervised stage: As mentioned before, the best time to introduce uncertainty to the full supervision phase is from the beginning. Uncertainty used here allows models in the network to obtain characteristics of the target area better under full supervision. We can see from Fig. 8. that the main effect of uncertainty in the full supervision stage is in the first 30 epochs. We can tell from the label that the uncertainty maps focus more on the lesion area. After adding uncertainty to the full supervision stage, the method can achieve improvement over co-training with an overall all performance increase of 2.23% in DSC and 0.57mm in HD.
In unsupervised stage: The best time at this stage is to leverage uncertainty after models have a certain un-derstanding of the lesion area's semantic features. After experiments per 10 epochs from 0, it is finally determined that the uncertainty will be leveraged after the 20th epoch. So, we visualize the uncertainty map from 0 epochs to the last in Fig. 9. We can see that the uncertainty maps are increasingly focusing on the lesion areas. Obviously, uncertainty always plays a role throughout the experiment. As the training progresses, the color of the lesion areas of the uncertainty maps becomes lighter, which means the predictions of the two models for the same unlabeled image are gradually approaching, but there is still a certain difference. Only introduce uncertainty to the unsupervised stage, we can achieve an improvement over co-training with an overall all performance increase of 1.39% in DSC and 0.75mm in HD.

Conclusion
In this paper, we propose a novel uncertainty-aware deep co-training method for three medical image segmentation tasks. Our approach uses uncertainty in both the supervised learning stage and the unsupervised learning stage in the deep co-training method. We use uncertainty obtained from Monte Carlo Sampling to guide the training process purposely. We validate our method on three challenging medical image datasets. The comparison with other semi-supervised methods confirms the effectiveness of our approach. Our uncertainty-aware co-training method achieves the performance 1.88% away from supervision with only 20% label on ACDC dataset, and an increase of 12.94% in terms of DSC over fully supervised method on SCGM dataset. And for the Spleen dataset, we get 0.89% more than the fully supervised approach in DSC. In future work, we will investigate the effect of different uncertainty estimation manners and apply our approach to other semi-supervised medical image segmentation tasks.