Space Engage: Collaborative Space Supervision for Contrastive-based Semi-Supervised Semantic Segmentation

The aim of S4 is to train a segmentation model with the semi-supervised setting (i.e., a few labeled images and a large number of unlabeled images) to classify each pixel in an entire image. The critical issue of S4 is how to leverage unlabeled images to train the model. Some methods [23, 25, 28, 33] based on GANs [14], adversarial training [34], and consistency regularization paradigm [36, 11, 38, 6, 54]. Meanwhile, self-training [26, 42, 47, 58, 51] is also a striking paradigm, which always generates pseudo-labels from model and retrains the model with the combined supervision of human annotations and pseudo-labels. One essential issue of self-training is the accuracy of pseudo-labels. Some methods [27, 31, 12, 40, 50] try to polish pseudo-labels and provide reliable guidance. Some methods [19, 45, 22, 15] focus on the class-imbalance problems in the dataset and try to alleviate the negative effect from class-biased pseudo-labels generated by the model pre-trained on imbalanced labeled images. We build our framework based on the self-training and additionally explore semantic information among different images.

Pixel-wise contrastive learning explores semantic relations not only in the individual image but also among different images. Different from instance-wise contrastive learning [17, 5, 2], pixel-wise contrastive learning [48, 53, 4, 56] project each pixel to the representation in representation space with the cooperation of encoder and representation head. Representations are then aggregated in their prototypes and are separated from each other in different classes. In semi-supervised settings, most methods [1, 30, 44, 46] use pseudo-labels based on logits to provide semantic information contrastive learning process during training on unlabeled images. Meanwhile, the confidence of logit is used as an indicator to involve the contrastive learning process, e.g., [30] uses the hard representations whose corresponding logit confidence is lower than a threshold to contrast for effective training. As opposed to the above methods, we use collaborative space supervision for contrastive learning on unlabeled images and use a new indicator to involve the contrastive learning progress.

Prototype-based learning has been widely studied in few-shot learning [39, 9, 29, 32] and unsupervised domain adaption [52, 24, 41, 37, 55]. Recently, it is restudied in semantic segmentation as known as a non-parametric prototype-based classifier [57]. Concretely, the classes in the dataset are presented by a set of non-learnable prototypes, and the dense semantic predictions are thus achieved by assigning the output features to its most similar prototype. Under semi-supervised settings, some methods [49] maintain the consistency between predictions from a linear predictor and a prototype-based predictor. The two predictors are followed by the encoder and project the features to logit space and representation space, respectively. In this work, we combine the semantic information in the logit and representation spaces to provide supervision in a collaborative way during semi-supervised learning.

The main idea of self-training is to pre-train a model on labeled images and use it to produce pseudo-labels as supervision for unlabeled images. A typical framework is the teacher-student framework [42], which consists of a student model and a teacher model. Both the student model and the teacher model are constructed by an encoder and a segmentation head. Parameters of the student model are optimized via Stochastic Gradient Decent (SGD) while parameters of the teacher model are updated by the Exponential Moving Average (EMA) of student model parameters. We denote the encoder and the segmentation head in the student model by $f(\cdot)$ and $g(\cdot)$ while denoting those in the teacher model by $f^{\prime}(\cdot)$ and $g^{\prime}(\cdot)$. The pseudo-labels $\hat{\bm{y}}_{i}^{u,lgt}$ are produced based on the teacher model’s output logits $\hat{\bm{p}}^{u}_{i}=g^{\prime}(f^{\prime}(\bm{x}^{u}_{i}))$ in logit space, formulated as:

$\displaystyle\begin{aligned} \hat{\bm{y}}_{i}^{u,lgt}=\bm{1}_{c}(\mathop{\arg\max}\limits_{c}\{\hat{\bm{p}}_{i,c}^{u}\}_{c\in C}),\end{aligned}$

(1)

where $\bm{1}_{c}(\cdot)$ denotes the one-hot encoding of class $c$.

In order to enhance the ability of the model itself to extract features, recent works [30, 44, 1] additionally employ pixel-wise contrastive learning and introduce a representation head to both teacher and student models. We denote the representation head in the student model as $h(\cdot)$ and that in the teacher model as $h^{\prime}(\cdot)$. The pixel $\bm{x}_{i}$ of the class $c$ are projected as representations $\bm{z}_{ci}$ in representation space by the cooperating of $f(\cdot)$ and $h(\cdot)$, i.e., $\bm{z}_{ci}=h(f(\bm{x}_{i}))$. And the representation $\bm{z}_{ci}$ is then aggregated to its class centroid (prototype) while separated from representations in different classes $\bm{z}_{\tilde{c}i}$ (negatives). The semantic guidance for contrastive learning is from the combination of ground truth $\bm{y}_{i}^{l}$ and pseudo-labels $\hat{\bm{y}}_{i}^{u,lgt}$ in logit space. Moreover, in order to emphasize the reliable and crucial aspects during unlabeled and contrastive learning, a sampling strategy is adopted to select valid pixels $\bm{x}_{i}$ according to their corresponding confidence, i.e. the student model’s output logits $\bm{p}_{i}$ after a Softmax operation.

To mitigate these limitations, we produce pseudo-labels from the representation space and combine them with pseudo-labels from the logit space to provide higher-quality supervision during unlabeled training. Meanwhile, we obtain a new indicator from the representation space for the more effective sampling strategy.

In this section, we detail the approach to obtain the pseudo-labels from the representation space. Meanwhile, we simultaneously access a new indicator for the sampling strategy in representation spaces, which provides a critical reference in the contrastive learning process.

Specifically, we first build a set of prototypes for each class and obtain the pseudo-labels by retrieving the nearest prototype for each representation. We calculate the centroid of all representations in the current class $c$ as the prototype $\bm{\rho}_{c}$, which is formulated as:

$\displaystyle\begin{aligned} \bm{\rho}_{c}=\frac{1}{N_{c}}\sum_{i}^{N_{c}}\bm{z}^{\prime}_{i},\end{aligned}$

(2)

where $N_{c}$ is the total number of representations of current class $c$ and $\bm{z}^{\prime}_{i}$ is the representation projected by the cooperation of the $f^{\prime}(\cdot)$ and $h^{\prime}(\cdot)$. Meanwhile, to include more representation information, we update the prototype along the sequential iterations with EMA as follows:

$\displaystyle\begin{aligned} \hat{\bm{\rho}}_{c}(t)=\alpha\hat{\bm{\rho}}_{c}(t-1)+(1-\alpha)\bm{\rho}_{c}(t),\end{aligned}$

(3)

where $\hat{\bm{\rho}}_{c}(t)$, $\hat{\bm{\rho}}_{c}(t-1)$ mean the current $t^{th}$ prototype and last $(t-1)^{th}$ prototype in iterations , $\bm{\rho}_{c}(t)$ means the prototype calculated by Eq. 2 in current iteration, and $\alpha$ is a hyper-parameter that controls the updating speed. Thus, the pseudo-label from the representation space is achieved by:

$\displaystyle\begin{aligned} \hat{\bm{y}}^{u,rep}_{i}=\bm{1}_{\hat{c}}(\hat{c}),with~{}\hat{c}=\mathop{\arg\max}\limits_{c}\{sim(\bm{z}^{\prime}_{i},\bm{\hat{\rho}}_{c}(t))\}_{c\in C},\end{aligned}$

(4)

where $sim(\cdot)$ is defined as the cosine similarity.

As for the indicator for the sampling strategy in the representation space, we use the Softmax function on the similarity among the representation and all prototypes, which is followed as:

$\displaystyle\begin{aligned} \hat{j}^{u,rep}_{i}=\frac{e^{sim(\bm{z_{ci}},\hat{\bm{\rho}}_{c}}(t))/\tau}{e^{sim(\bm{z_{ci}},\hat{\bm{\rho}}_{c}(t))/\tau}+\sum_{\tilde{c}\in\tilde{C}}e^{sim(\bm{z_{ci}},\hat{\bm{\rho}}_{\tilde{c}}(t))/\tau}},\end{aligned}$

(5)

where $\hat{\bm{\rho}}_{\tilde{c}}(t)$ means the prototype with different classes from $\bm{z}_{ci}$ and $\tau$ is a hyper-parameter. Different from using confidence from logit space as the indicator to involve representation learning [30, 44], the Softmax similarity directly helps the model to discover the confusion between representations and their prototypes and focus on them during the subsequent training.

With the pseudo-labels in two spaces, we propose two pseudo-labeling strategies to strengthen the collaboration between two spaces and obtain more reliable pseudo-labels.

• •

  Mix pseudo-labeling. To mitigate the negative effects of inherent noise from both two spaces during pseudo-labeling, we adopt the mix pseudo-labeling strategy that only considers the mutually agreeable pseudo-labels between the two spaces. Specifically, we define the set of final pseudo-labels as $\hat{Y}^{u}=\hat{Y}^{u,lgt}\cap\hat{Y}^{u,rep}$, where $\hat{\bm{y}}_{i}^{u,lgt}\in\hat{Y}^{u,lgt}$ and $\hat{\bm{y}}_{i}^{u,rep}\in\hat{Y}^{u,rep}$.

• •

  Cross pseudo-labeling. Inspired by recent researches [6, 36] that maintain consistency among the predictions of the same image across different models or decoders in different views, we propose a cross pseudo-labeling strategy that leverages pseudo-labels from one space to supervise the other. Specifically, we use pseudo-labels $\hat{\bm{y}}_{i}^{u,rep}$ to supervise the logit space, and vice versa.

The strengths of using pseudo-labels from two spaces in the collaborative way are twofold: 1) obtaining more reliable supervision during unlabeled training, and 2) enabling the strengths of learning in different spaces to complement each other. Since the learning in different feature spaces concentrates on different parts of features, i.e., the logit space mainly focuses on the most discriminative part of features while the representation space treats all parts of features equally, the performance of pseudo-labels from two spaces varies across different classes or regions of images. Therefore, our collaborative pseudo-labeling strategies exchange knowledge between two spaces and provide higher-quality supervision during unlabeled training. The experimental proof is in Sec. 5.1.

As for indicators, we use confidence as the indicator $\hat{j}_{i}^{u,lgt}$ for learning logit space and Softmax similarity as the indicator $\hat{j}_{i}^{u,rep}$ for learning representation space. We argue that the confusing parts of learning in both two spaces are varied due to the different parts of features being concentrated in each space. Therefore, this strategy allows the learning in different spaces to focus on their own confusing parts, which can be more effective than using a single indicator when mining confusing parts for both spaces in the training process. The experimental proof is in Sec. 5.2.

With the indicators $\hat{j}_{i}^{u,lgt}$ and $\hat{j}_{i}^{u,rep}$, we adopt some threshold sampling strategies. In logit space, we set a threshold $\delta_{u}$ during unlabeled learning and logits $\hat{\bm{p}}_{i}^{u}$ whose indicator $\hat{j}_{i}^{u,lgt}$ is higher than $\delta_{u}$ will be regarded as the valid logits in logit space. In representation space, our sample strategy can be divided into three parts: 1) Valid Sampling Strategy. Similar to the sampling strategy in logit space, a threshold $\delta_{w}$ is used to sample representations whose indicator $\hat{j}_{i}^{u,rep}$ is higher than $\delta_{w}$. 2) Hard Sampling Strategy. We adopt the hard sampling strategy for tilting to train those confusing representations. Specifically, we set a threshold $\delta_{s}$ to sample representations whose indicator $\hat{j}_{i}^{u,rep}$ is lower than $\delta_{s}$. 3) Negative Sampling Strategy. We sample negatives according to the similarity between the prototype of current class $c$ and other prototypes. Concretely, the negatives are more likely to be sampled if its prototype is more similar to the prototype of the current class.

Cooperated with the ground truth $\bm{y}^{l}_{i}$, pseudo-labels $\bm{y}^{u}_{i}$ produced from two spaces in a collaborative way, and different sampling strategies in two spaces, the total learning object is composed with a supervised loss $\mathcal{L}_{s}$, an unsupervised loss $\mathcal{L}_{u}$, and a contrastive loss $\mathcal{L}_{c}$ as follows:

$\displaystyle\mathcal{L}=\mathcal{L}_{s}+\mathcal{L}_{u}+\lambda_{c}\mathcal{L}_{c},$

(6)

where $\lambda_{c}$ is used to tune the contribution between logit space and representation space. Specifically, $\mathcal{L}_{s}$ and $\mathcal{L}_{u}$ are constructed by the Cross-Entropy (CE) $\ell_{ce}$ and can be formulated as:

$\displaystyle\mathcal{L}_{s}=\frac{1}{\lvert\mathcal{B}_{l}\rvert}\sum_{\bm{x}_{i}^{l}\in\mathcal{B}_{l}}\ell_{ce}(\bm{p}_{i}^{l},\bm{y}^{l}_{i}),$

(7)

$\displaystyle\mathcal{L}_{u}=\frac{1}{\lvert\hat{\mathcal{B}}_{u}\rvert}\sum_{\bm{x}_{i}^{u}\in\hat{\mathcal{B}}_{u}}\ell_{ce}(\bm{p}_{i}^{u},\hat{\bm{y}}^{u}_{i}),$

(8)

where $\mathcal{B}_{l}$ denotes the labeled images in a mini-batch and $\hat{\mathcal{B}}_{u}$ is the subsets that sampled from unlabeled images in a mini-batch according to the sampling strategy. Meanwhile, the contrastive loss $\mathcal{L}_{c}$ is formulated as:

$\displaystyle\begin{aligned} \mathcal{L}_{c}=&-\frac{1}{|C|\times|\hat{\mathcal{Z}_{c}}|}\sum_{c\in C}\sum_{\bm{z}_{ci}\in\hat{\mathcal{Z}}_{c}}\\ &log[\frac{e^{s(\bm{z}_{ci},\hat{\bm{\rho}}_{c})/\tau}}{e^{s(\bm{z}_{ci},\hat{\bm{\rho}}_{c}(t)))/\tau}+\sum_{\tilde{c}\in\tilde{C}}\sum_{\bm{z}_{\tilde{c}i}\in\hat{\mathcal{Z}}_{\tilde{c}}}e^{s(\bm{z}_{ci},\bm{z}_{\tilde{c}i})/\tau}}],\end{aligned}$

(9)

where $\hat{\mathcal{Z}_{c}}$ is the subset sampled from the set of the representations which belong to class $c$ according to the sampling strategy, $\hat{\mathcal{Z}}_{\tilde{c}}$ is the subset sampled from the set of the representations which bot belong to class $c$, $\tilde{C}$ denotes the subset of $C$ with class $c$ removed, and the supervision information comes from the final pseudo-labels $\hat{\bm{y}}^{u}_{i}$ after the pseudo-labeling strategies. The whole framework is shown in Fig. 2.

In this subsection, we first reproduce three baselines: MT [42], CutMix [13], and a typical contrastive-based method with only logit space pseudo-labels and indicators (Baseline) on classic VOC train set. Meanwhile, we make the comparison of our method with mix (CSS (mix)) and cross (CSS (crs.)) pseudo-labeling strategy on blender VOC train set and Cityscapes train set with following recent SOTA S4 methods: CCT [36], CPS [6], U²PL [44], ST++ [50], PRCL [46], PCR [49], and PSMT [31]. Since the data split will dramatically affect the performance in S4, i.e., choosing labeled images plays an important role in the final results, we conduct experiments with three different data splits and report the mean value and standard deviation (blue color numbers). Since the mix pseudo-labeling strategy has better performance, we only use CSS (mix) when compared with SOTAs. Meanwhile, we use ResNet-101 with deep stem blocks as our network structure when compared with SOTAs. Since there is no uniform data split, we use the data splits in U²PL [44]. We use the OHEM loss when training Cityscapes, following [6].

Meanwhile, we also visualize the pseudo-labels obtained from logit space (lgt.) and representation space (rep.) in Fig. 3. Fig. 3 (c) and (e) show the masks for pseudo-labels produced by the sampling strategies. In particular, the white color represents the valid pixels used during unlabeled learning, while the black color indicates the discarded pixels. Fig. 3 (d) and (f) are the pseudo-labels we obtained from two spaces. The figure clearly illustrates the differences between pseudo-labels produced by different spaces. For example, the parts of the instance edge in pseudo-labels are usually discarded since they are challenging for learning in logit space. However, pseudo-labels from representation space will easily tackle this problem (shown in the first row). The pseudo-labels from representation space are more inaccurate in some complex scenes, which be resolved by combining pseudo-labels from logit space (shown in the second row).

We mainly attribute the differences in pseudo-label to the differing concentrations of learning in the two spaces. Specifically, learning in the logit space primarily emphasizes the most discriminative part of features, whereas that in the representation space treats each part of features equally. As a result, learning in the logit space may overlook minor feature differences, leading to sub-optimal performance in predicting instance edges and distinguishing between similar classes (e.g., chair and sofa). Conversely, learning in the representation space produces balanced performance across all image parts and classes. However, this can lead to erroneous predictions for classes with high intra-class variance (e.g., background). By leveraging pseudo-labels from the two spaces, we capitalize on the strengths of the learning in each space and enhance the knowledge exchange between the two spaces.

Fig. 5 visualizes the similarity and confidence of an image in both logit (lgt.) and representation (rep.) spaces, indicating the varying levels of confusion in the same region when learning in different spaces.

We also attribute it to the different concentrations of learning two spaces, i.e., the confusing region in one space can be more readily addressed in the other space.

Thus, it is inappropriate to choose confidence as the indicator to involve representation learning, e.g., sampling more reliable or hard samples with threshold and indicator for better learning. In contrast, our indicator directly employs similarity between the representation and prototype of its class, which directly reflects the confusion level in representation learning. It is more accurate to use similarity as the indicator to sample hard and critical samples in representation learning.

In this section, we conduct experiments to introduce our components in CSS step by step, with results shown in Tab. 7. Our baseline is the conventional contrastive-based S4, achieving mIoU of $67.11\%$ on 92 labels and $70.32\%$ on 183 labels. Mix and cross means the pseudo-labels are from the mix pseudo-labeling strategy and cross pseudo-labeling strategy while the indicator is still the confidence in two spaces. Ind means we use the different indicators in different spaces while the pseudo-labels are from logit space. The last two rows represent our proposed two pseudo-labeling strategies with indicators from two spaces. As the result, the mix pseudo-labeling strategy with different indicators boosts model performance by $1.30\%$ and $2.42\%$ while the cross pseudo-labeling strategy with different indicators increases the performance by $0.74\%$ and $1.66\%$.

Fig. 6 shows the qualitative results of different methods on PASCAL VOC 2012 with 92 labeled images. Baseline means the conventional contrastive-based method. Compared with the original self-training methods (CutMix), thanks to introducing pixel-wise contrastive learning, the baseline, and our method perform better in some ambiguous regions. Furthermore, benefiting from the supervision of two spaces and different indicators in different spaces, our method outperforms the baseline.