Pseudo-label Alignment for Semi-supervised Instance Segmentation

The goal of PAIS is to better leverage unlabeled images with a limited number of pixel-labeled images to boost the performance of semi-supervised instance segmentation. The PAIS framework consists of three key steps: (1) pseudo-label generation, (2) dynamic pseudo-label alignment, and (3) end-to-end model training. In the pseudo-label generation step, we introduce a mask scoring branch [17] that predicts mask IoUs as an additional metric along with classification scores to assess the quality of pseudo-labels. In the dynamic aligning step, we re-weight the loss terms based on the quality of different pseudo-labels. Finally, the teacher and student models are trained using exponential moving average (EMA) [44]. The overall framework of PAIS is illustrated in Fig. 2, and we introduce the key steps in detail as follows. It is important to note that the PAIS framework can be applied to any box-dependent or box-free framework for enhanced exploitation of pseudo-labels in semi-supervised instance segmentation. We instantiate two examples by Mask-RCNN [13] and K-Net [53] in this paper, but do not restrict to them.

In the pseudo-label generation, we apply weak data augmentation to the unlabeled images and input them into a teacher model. The weak data augmentation includes scaling, horizontal flipping, and other augmentation operations that do not alter the image’s content. The teacher model produces a set of pseudo-labels, including masks, boxes, and classification scores for each input image. In box-free pipelines, the box predictions are optional.

As depicted in Fig. 1(a)(c), using the classification score alone is inadequate to measure the quality of predicted masks. To address this, we incorporate a mask scoring branch into the pipeline to predict mask Intersection over Union (IoU) for evaluating mask quality. As shown in Fig. 1(b), the predicted mask IoUs can be used to measure mask quality effectively, which has also been verified in MS-RCNN [17]. After generating pseudo-labels in the previous step, the predictions are filtered using two thresholds: a classification threshold $\tau_{cls}$ and a mask IoU threshold $\tau_{iou}$. The resulting set of pseudo-labels includes $N$ elements, each containing a mask $\boldsymbol{m}_{i}\in\mathbb{R}^{H\times W}$, a bounding box $\boldsymbol{b}_{i}\in\mathbb{R}^{4}$, a classification score $\boldsymbol{c}_{i}\in\mathbb{R}^{L}$ (with an additional background category), and a mask IoU score ${s}_{i}\in\mathbb{R}$. Note that $H,W$ denote the mask resolution, and $L$ is the number of classes.

Although good pseudo-labels can be obtained by setting high thresholds for classification scores and mask IoUs, a large number of predictions with misaligned mask and classification qualities are discarded (as illustrated in Fig.1(c)(d)). These misaligned predictions can be useful since the amount of labeled data is limited in the semi-supervised setting. Fig.2 shows examples of pseudo-labels that can be included in the training process. If the misaligned pseudo-labels are directly used in semi-supervised learning, the incorrect predictions on classes or masks can introduce significant noise. To reduce the noise, we propose a dynamic aligning loss (DALoss).

Vanilla Loss for PAIS. In semi-supervised instance segmentation, the loss function can be decomposed into two terms for labeled and unlabeled images, as:

where $\lambda_{l}$ and $\lambda_{u}$ are the hyper-parameters for balancing the loss terms.

The loss for labeled images can be defined using the functions commonly employed in instance segmentation, augmented with an additional binary cross-entropy term for regressing mask IoUs. Given the pseudo-labels from the teacher model and the predictions from the student model, we formulate the loss for unlabeled images as:

where $\mathcal{L}_{b}$ includes the box IoU loss and the L1 loss, $\mathcal{L}_{c}$ denotes the cross-entropy loss, $\mathcal{L}_{m}$ is the dice loss. In Eq. 2, $\widetilde{\boldsymbol{c}}_{i},\widetilde{\boldsymbol{b}}_{i},\widetilde{\boldsymbol{m}}_{i}$ represent predictions from the student model. The pseudo-scores $\boldsymbol{c}_{\mu(i)}$ for classification are converted to one-hot vectors $\boldsymbol{c}_{\mu(i)}^{o}$ for the category with the highest score. The pseudo-masks $\boldsymbol{m}_{\mu(i)}$ are activated by sigmoid function and discretized into binary value as $\boldsymbol{m}_{\mu(i)}^{b}$. The one-hot vector for the background class is denoted by $\boldsymbol{c}_{neg}^{o}$. The label indexing functions, $\nu(i)$ and $\mu(i)$, match predictions with pseudo-labels for training, which are defined in terms of different pseudo-label assigning strategies. In our implementations, the one-to-many assignment defines the label indexing functions $\nu(i)=i$ and $\mu(i)$ as:

Instead, the one-to-one assignment defines the label indicting functions $\mu(i)=i$, and finds the optimal $\nu^{*}(\cdot)$ via:

Note that the loss terms for regressing boxes are optional in box-free instance segmentation frameworks.

Dynamic Aligning Loss for PAIS. To optimize the model using pixel-level pseudo-labels, we propose to replace Eq. 2 with DALoss. Since the teacher model provides ideal metrics for measuring the quality of classification scores and mask IoUs, we propose to adjust the weight of loss terms based on the quality of the pseudo-labels. This is achieved by using the following equation:

where $c_{\mu(i)}^{h}$ denotes the highest classification score, and $s_{\mu(i)}$ denotes the mask IoU from the $\mu(i)$-th pseudo-label, $\alpha,\beta$ are the hyper-parameters. Specifically, Eq. 5 adjusts the weights for the pseudo-labels conditioned on their qualities, i.e., dynamically dependent on the input images. For instance, for a pseudo-label with a low classification score and high mask IoU, DALoss encourages the segmentation loss while constraining the classification loss.

Inspired by [44], we employ EMA with the strong-weak data augmentation for PAIS. Specifically, unlabeled images undergo both strong and weak data augmentations and are then fed into the student and teacher models, respectively. The student model is trained to produce consistent results with the pseudo-labels, and the teacher model is updated by EMA. The training pipeline for PAIS is presented in Alg. 1.

We conducted extensive experiments on the COCO [25] and Cityscapes [6] datasets to study the proposed PAIS. The COCO dataset consists of 118k images with 80-class instance labels, as well as 123k unlabeled images. The Cityscapes dataset contains urban street-view scenes and has 8 instance categories in 2.9k training images and 0.5k validation images. For the COCO dataset, we randomly sampled 1%, 5%, and 10% of the images from the train2017 split as labeled data and treated the rest as unlabeled data following common settings. In addition, we also used the full COCO train2017 as labeled data and incorporated the 123k unlabeled data from COCO unlabel2017 to train the PAIS models. For the Cityscapes dataset, we randomly sampled 5%, 10%, 20%, and 30% of the images from the training set as labeled data and treated the remaining as unlabeled data following the common settings. We evaluated the PAIS models on the validation sets of the COCO and Cityscapes datasets, and reported standard COCO metrics including AP, AP${50}$, AP${75}$ (averaged over IoU thresholds), and AP_S, AP_M, AP_L (AP for instances of different scales).

We provide two examples of implementing PAIS with K-Net [53] and Mask-RCNN [13]. The models are trained using AdamW with a learning rate of 0.0001 for K-Net, and SGD with a learning rate of 0.01 for Mask-RCNN. The hyper-parameters $\lambda_{l}$ and $\lambda_{u}$, which balance the loss terms for labeled and unlabeled images, are set to 1.0 and 0.3 for K-Net, and 1.0 and 1.5 for Mask-RCNN. We set the thresholds $\tau_{cls}$ and $\tau_{iou}$ experimentally as 0.35 and 0.30, respectively. For bipartite matching loss, we use the same hyperparameters as in [53]. The loss balancing parameters for box, class, and mask are set as $\lambda_{b}$=2.0, $\lambda_{c}$=4.0, and $\lambda_{m}$=1.0, respectively. We train the models on 4 GPUs with 4 images per GPU (1 labeled and 3 unlabeled images) for 220k iterations, unless otherwise specified. The teacher model is updated via EMA with a momentum of 0.999. We use ResNet50 [14] as the backbone for these models.

Comparison to state-of-the-art semi-supervised instance segmentation frameworks. In Tab. 1, we compare the performance of models trained with PAIS to the state-of-the-art semi-supervised instance segmentation frameworks on the COCO dataset. The results demonstrate that both K-Net and Mask-RCNN trained with PAIS surpass the previous methods DD [35] and Noisy Boundaries [47] by a large margin, especially when the number of labeled data is very limited (with only 1% or 5% labeled images). Specifically, when using 1% labeled COCO images, PAIS with K-Net achieves 19.9 mask mAP, which is 12.2 points higher than Noisy Boundaries. Interestingly, the proposed PAIS brings about better performance for Mask-RCNN when the percentage of labeled data is 1% or 5%, even though it originally has a inferior performance in fully-supervised instance segmentation compared to K-Net. To explain, the bounding boxes may provide better optimization for Mask-RCNN models when the labeled data is limited. Finally, when using 10% and 100% labeled images, PAIS with K-Net outperforms PAIS with Mask-RCNN. In Tab. 2, we compare PAIS with state-of-the-art methods on the Cityscape dataset, in which PAIS also achieves better performance than the predominant models. Specifically, we report NoisyBoundaries [47] w/o FocalLoss in Tab. 2, as we do not apply FocalLoss in PAIS. This ensures a fair and consistent comparison. Furthermore, we also add FocalLoss to PAIS on 10% Cityscapes, which achieves 25.1%, surpassing 23.7% of NoisyBoundaries w/ FocalLoss.

The more significant performance improvement on the COCO dataset validates the effectiveness of our method for solving the noisy pseudo-label problem. The COCO dataset has 80 instance categories, while the Cityspaces dataset only has 8 instance categories. This implies that the COCO dataset can provide more diverse and informative pseudo-labels, which matches our goal to utilize noisy pseudo-labels for semi-supervised learning.

Results with an extremely limited number of labeled images. To demonstrate the effectiveness of PAIS, we conduct experiments with extremely limited numbers of labeled images, as shown in Table 3. Specifically, we compare the performance of various models trained with randomly sampled 1%, 5%, and 10% labeled images. First, we train supervised models, K-Net (supervised) and Mask-RCNN (supervised), with the limited labeled images. Second, we train the same supervised models with the same data augmentation used in the semi-supervised setting, denoted as K-Net^† (supervised) and Mask-RCNN^† (supervised). Third, we train PAIS models without DALoss, denoted as PAIS w/o DALoss on K-Net and Mask-RCNN. Lastly, we train the PAIS models with DALoss, denoted as PAIS on K-Net and Mask-RCNN. All models are trained three times, and the reported results are averaged.

Based on the results of K-Net (supervised) and Mask-RCNN (supervised), it can be observed that fully-supervised models perform poorly when the number of labeled images is limited. The results of K-Net^† (supervised) and Mask-RCNN^† (supervised) show slight improvement with weak data augmentation from the semi-supervised setting. Interestingly, while K-Net outperforms Mask-RCNN in the fully-supervised setting, their performance is similar when the number of labeled images is limited, as indicated in the table. Comparing the performance of K-Net (PAIS w/o DALoss) and Mask-RCNN (PAIS w/o DALoss) with that of the supervised setting reveals significant improvement. By introducing the dynamically re-weighting process via DALoss, the performance is further improved. For instance, with only 1% labeled images, the mAP of K-Net and Mask-RCNN improved from 11.63 and 11.54 to 19.78 and 21.12, respectively, resulting in approximately +8.15 and +9.58 points improvement to the performance. Additionally, the improvement over AP_0.5 and AP_0.75 suggests that DALoss considers masks with moderate quality during training, which further benefits semi-supervised learning.

Results with abundant labeled images. In Tab. 4, we investigate the performance of semi-supervised learning when abundant labeled data is available. Specifically, we use the entire COCO train2017 dataset as labeled data and COCO unlabel2017 dataset as unlabeled data to train the models. The results show that the semi-supervised learning approach also leads to a performance gain. For instance, K-Net (PAIS) achieves a performance gain of approximately 1.0 point on mAP from semi-supervised learning. Additionally, the proposed DALoss consistently improves the performance of both box-dependent and box-free instance segmentation frameworks.

In our ablation study, we investigate several aspects that can impact the performance of PAIS, including the utilization of various loss terms, the setting of hyper-parameters, the threshold values, the varying ratios of labeled and unlabeled images, and the convergence times. We perform the ablation studies on the COCO dataset using K-Net and Mask-RCNN, which are trained via PAIS under the setting of 10% labeled images.

Effectiveness of different terms in DALoss. Tab. 5 shows the efficacy of the different components in DALoss, which suggests that: (1) DALoss yields an improvement of 1.8 points to the model. Specifically, the mAP increases from 45.5 to 48.3 on AP_L, indicating a significant improvement for large objects. (2) The aligning weights for classification loss alone can also provide a marginal performance gain for the model. (3) Simply adding a mask IoU branch to the model does not enhance the overall performance. However, when incorporating the aligning weights for the segmentation loss, the mAP is increased. (4) Interestingly, the mask IoU branch can help improve the AP for objects of different scales. This is due to the mask IoU branch helping to select good pseudo-labels for training.

Hyper-parameters in DALoss. We investigated the effect of hyper-parameters, $\alpha$ and $\beta$, used in DALoss. To give equal weight to both classification and segmentation, we set $\alpha=\beta$. The results are shown in Tab. 6, indicating that larger values of $\alpha$ and $\beta$ lead to better performance. We also analyzed the functions of different hyper-parameters used to adjust the input score, as shown in Fig. 3. The figures illustrate that increasing the hyper-parameter enlarges the difference between high- and low-quality masks or classes, adjusting the input score more appropriately based on their quality. However, when the hyper-parameter is set too large, the low-quality masks or classes are significantly constrained, potentially affecting the generalization of the model as noise is removed entirely during training.

Model convergence speed. In Fig. 4, we analyze the convergence speed of the models. The results demonstrate that DALoss can expedite the model training in terms of convergence rate. For instance, Mask-RCNN (PAIS) achieves a convergence speed that is approximately 2 times faster than that of Mask-RCNN (PAIS w/o DALoss). Moreover, by comparing the results between Mask-RCNN (PAIS) and K-Net (PAIS), we observe that DALoss may be particularly effective for box-dependent instance segmentation frameworks, which facilitates rapid convergence.

Different thresholds for classification scores and mask IoUs. We observed from Tab. 7 and Tab. 8 that increasing the thresholds for classification score and mask IoU leads to a decrease in performance. This suggests that DALoss is able to make use of misaligned pseudo-labels instead of simply filtering them out. By allowing the model to learn from these potentially noisy labels, it is able to better handle situations where the alignment between labeled and unlabeled data is not perfect. This may result in improved generalization to new, unseen data, as the model has learned to adapt to the presence of noisy labels.

Different ratios of labeled and unlabeled images per batch. Tab. 9 shows that increasing the ratio of unlabeled images in a batch can improve the model’s performance. The results indicate that performance saturates when the ratio is set to 1:4, suggesting that adding too many unlabeled images may lead to diminishing returns.

Visualizations. We visualize the outputs of different models, including K-Net (supervised), K-Net (PAIS w/o DALoss), K-Net (PAIS), Mask-RCNN (supervised), Mask-RCNN (PAIS w/o DALoss), and Mask-RCNN (PAIS), in Fig. 5. The results show that: (1) PAIS can improve the recall for most instances in the supervised model, but the quality of the predictions is not guaranteed. (2) The use of DALoss in PAIS helps to improve the quality of predictions in terms of mask and classification.

Influence of imbalanced classes. In Fig. 6, we investigate the effect of imbalanced classes on PAIS. We plot the number of labels obtained from the ground truth and the teacher models for Mask-RCNN and K-Net with PAIS at different iterations, namely 32k, 64k, 120k, and 180k. The results indicate that the predicted pseudo-labels gradually conform to the distribution of the imbalanced labels.