Cross-Image Relational Knowledge Distillation for Semantic Segmentation

Introduction

Semantic segmentation is a crucial and challenging task in computer vision. It aims to classify each pixel in the input image with an individual category label. The applications of segmentation often focus on autonomous driving, virtual reality and robots. Although popular state-of-the-art segmentation networks, such as DeepLab , PSPNet and OCRNet , achieve remarkable performance, they often need high computational costs. This weakness makes them difficult to be deployed for real-world scenarios over resource-limited mobile devices. Therefore, a series of lightweight segmentation networks are proposed, such as ESPet , ICNet and BiSeNet . Moreover, model compression is also an alternative field to pursue compact networks, mainly divided into quantization , pruning and knowledge distillation (KD) .

This paper investigates KD to improve the performance of a compact student network under the guidance of a high-capacity teacher network for semantic segmentation. A broad range of KD approaches have been well studied but mostly for image classification tasks. Unlike image-level recognition, the segmentation task aims at dense pixel predictions, which is more challenging. Previous researches have found that directly utilizing classification-based KD methods to deal with dense prediction tasks may not achieve desirable performance. This is because strictly aligning the coarse feature maps between the teacher and student networks may lead to negative constraints and ignore the structured context among pixels.

Recent works attempt to propose specialized KD methods for semantic segmentation. Most focus on mining correlations or dependencies among spatial pixel locations because segmentation needs a structured output. Typical knowledge can be local pixel affinity , global pairwise relations and intra-class pixel variation . Such methods often perform better than the traditional point-wise alignment in capturing structured spatial knowledge. More recently, Shu et al. revealed that each channel represents a category-specific mask and thus proposed Channel-Wise KD (CWD) . CWD achieves state-of-the-art distillation performance and demonstrates the importance of channel-level information for dense prediction tasks. However, previous segmentation KD methods often guide a student to mimic the teacher’s structured information generated from individual data samples. They ignore cross-image semantic relations among pixels for knowledge transfer, as shown in Fig. 8.

Based on this motivation, we propose Cross-Image Relational Knowledge Distillation (CIRKD) for semantic segmentation. The core idea is to construct global pixel relations across the whole training images as meaningful knowledge. A good pre-trained teacher network could often generate a well-structured pixel embedding space and capture better pixel correlations than a student network. Based on this property, we transfer such pixel relations from teacher to student. Specifically, we propose pixel-to-pixel distillation and pixel-to-region distillation to fully exploit structured relations across various images. The former aims to transfer similarity distributions among pixel embeddings. The latter focuses on transferring pixel-to-region similarity distributions complementary to the former. The region embedding is generated by averagely pooling pixel embeddings from the same class and represents that class’s feature center. The pixel-to-region relations indicate the relative similarities between pixels and class-wise prototypes.

A naive way for constructing cross-image relations is to derive embeddings from the current mini-batch. However, the batch size of the segmentation task is often small, limiting the network to capture broader pixel dependencies. Motivated by previous self-supervised learning , we introduce a pixel queue and a region queue in the memory bank to store abundant embeddings for modelling long-range pixel relations. The embeddings in queues are consistent during the distillation process, since they are generated from the pre-trained and frozen teacher network. We regard the teacher and student pixel embeddings from the current mini-batch as anchors. We randomly sample contrastive embeddings from the queues to model pixel-to-pixel as well as pixel-to-region similarity distributions. Then we align such soft relations via KL-divergence from the student to teacher.

CIRKD guides the student network to learn the global property of relative pixel structures across training images from the teacher, further improving the segmentation performance. We evaluate our method over popular DeepLabV3 and PSPNet architectures on three segmentation benchmark datasets: Cityscapes , CamVid and Pascal VOC . Experimental results indicate that CIRKD outperforms other state-of-the-art distillation approaches, demonstrating the value of transferring global pixel relationships in semantic segmentation.

The main contributions are summarized as follows:

We propose cross-image relational KD to transfer global pixel relationships. We may be the first to build pixel dependencies across global images for segmentation KD.

We propose pixel-to-pixel and pixel-to-region distillation with the memory bank mechanism to fully explore structured relations for transfer.

Our CIRKD achieves the best distillation performance among state-of-the-art methods on the public segmentation datasets.

Related Work

Semantic Segmentation. Fully Convolutional Networks (FCN) creates a seminal paradigm for end-to-end dense feature learning for semantic segmentation. Since contextual pixel dependencies are essential for segmentation performance , capturing long-range relationships becomes a critical topic. DeepLab applies atrous convolution to enlarge the receptive field for learning broader context. DeepLabV3 assembles convolution blocks with various atrous rates in parallel to capture multi-scale contexts. PSPNet proposes a pyramid pooling module to exploit different-region-based context aggregation. RefineNet preserves high-resolution predictions by long-range residual connections for the down-sampling process. More recently, SegFormer utilizes a structured Transformer encoder to model global context information. However, such high-performance segmentation networks with expensive computational costs are difficult to be deployed over resource-limited mobile devices.

Efficient segmentation networks attract wide attention due to the need for real-time inference. Most works attempt to design lightweight networks with cheap operations. ENet is equipped with early downsampling, small decoder size and filter factorization. ESPNet factorizes the standard convolution into the spatial pyramid of dilated convolution. ICNet builds a cascade structure to balance the efficiency between low-resolution and high-resolution features. BiSeNet combines a spatial path and a context path to process features efficiently. Beyond designing a segmentation framework, lightweight backbone networks , e.g. MobileNet and ShuffleNet , can also implement acceleration.

Knowledge Distillation. The core idea of KD is to transfer meaningful knowledge from a cumbersome teacher into a smaller and faster student. Most current KD methods deal with image classification networks, mainly divided into probability-based, feature-based and relation-based approaches. Probability-based KD transfers class probabilities produced from the teacher as soft labels to supervise the student. Feature-based KD focuses on intermediate feature maps or their refined information as knowledge. Relation-based KD aligns correlations or dependencies among multiple instances between the student and teacher networks. Our CIRKD is related to SEED that both of them are contrastive distillation manners with a shared memory bank. However, these image-level KD methods are often unsuitable for pixel-wise semantic segmentation .

Recent KD methods for semantic segmentation often encode contextual pixel affinity as knowledge. Xie et al. align local similarity maps constructed from 8 neighbourhood pixels between the student and teacher networks. He et al. transfer non-local pairwise affinity maps with an autoencoder to minimize the discrepancy of features. Liu et al. perform a pairwise similarity distillation among pixels and an adversarial distillation of score maps. Wang et al. distill intra-class feature variation to learn more robust relations with class-wise prototypes. Beyond spatial distillation, Shu et al. propose channel-wise distillation to guide the student to mimic the teacher’s semantic masks along the channel dimension. Though achieving desirable performance, these approaches only consider pixel dependencies within an individual image, ignoring global pixel relations across various images.

Methodology

Here, $CE$ denotes the cross-entropy loss, $\sigma$ denotes the softmax function and $y_{h,w}$ denotes the ground-truth label of the $(h,w)$-th pixel.

Pixel-wise Class Probability Distillation. Motivated by Hinton’s KD , a direct method is to align the class probability distribution of each pixel from the student to the teacher. The formulation is expressed as:

Here, $\sigma(\mathbf{Z}^{s}_{h,w}/T)$ and $\sigma(\mathbf{Z}^{t}_{h,w}/T)$ represent the soft class probabilities of the $(h,w)$-th pixel produced from the student and teacher, respectively. $KL$ denotes the Kullback-Leibler divergence, and $T$ is a temperature. Following previous works , $T=1$ is good enough.

2 Cross-Image Relational Knowledge Distillaton

Motivation. Although the training objectives of $L_{task}$ and $L_{kd}$ are widely used in semantic segmentation, they only deal with pixel-wise predictions independently but neglect semantic relations between pixels. Some segmentation KD methods attempt to capture spatial relational knowledge by modelling pixel affinity. Nevertheless, these KD methods only construct the relationships among pixels within a single image, regardless of the semantic dependencies among pixels across global images. This paper demonstrates that cross-image relational knowledge is also valuable for conducting teacher-student-based KD.

Our CIRKD makes use of pixel embeddings beyond a single image. Inspired by the recent memory-based contrastive learning , we may retrieve pixel embeddings of other images from the current mini-batch or an online memory bank. This paper considers both of two manners to model relationships among pixels, the details of which are shown as follows.

We guide the pair-wise similarity matrix of $\mathbf{S}^{s}_{ij}$ produced from the student to align that of $\mathbf{S}^{t}_{ij}$ produced from the teacher. The distillation process is formulated as:

Here, $\mathbf{S}_{ij|a,:}$ denotes the $a$-th row vector of $\mathbf{S}_{ij}$. We normalize each row similarity distribution of $\mathbf{S}_{ij}$ to a probability distribution with a temperature $\tau$ by softmax function $\sigma$. The magnitude gaps would be removed between the student and teacher networks due to the softmax normalization. $KL$ is used to align each row-wise probability distribution. We perform pixel-to-pixel distillation every two of $N$ images:

We show the illustration of mini-batch-based pixel-to-pixel distillation in the supplement.

2.2 Memory-based Pixel-to-Pixel Distillation

Although mini-batch-based distillation could capture cross-image relations to some extent, it is difficult to model dependencies among pixels from global images, since the batch size per GPU of segmentation task is often small, e.g. 1 or 2. To address this problem, we introduce an online pixel queue that can store massive pixel embeddings in the memory bank generated from the past mini-batches. It allows us to retrieve abundant embeddings efficiently. The usage of memory bank dates back to self-supervised learning . This is because a large number of negative samples are pivotal for unsupervised contrastive learning, and the mini-batch size limits available contrastive samples.

The teacher network often shows a better pixel similarity matrix than the student. We force the student’s $\mathbf{P}^{s}$ to mimic the teacher’s $\mathbf{P}^{t}$ for penalizing the difference. Similar to Section 3.2.1, we apply softmax normalization on each row distribution of $\mathbf{P}^{s}$ and $\mathbf{P}^{t}$ and perform pixel-to-pixel distillation via KL-divergence loss. It is formulated as follows:

After each iteration, we push $V$ teacher pixel embeddings per class into the pixel queue $\mathcal{Q}_{p}$. Because the teacher is pre-trained and frozen, it can provide consistent feature embeddings during the distillation process. Therefore, we can naturally avoid the inconsistent problem between the anchor and dequeued features appeared in previous contrastive learning .

2.3 Memory-based Pixel-to-Region Distillation

Discrete pixel embeddings may not fully capture image content. Thus we introduce an online region queue that can store massive more representative region embeddings in the memory bank. Beyond pixel-to-pixel distillation, we further construct pixel-to-region distillation to model the relations between pixels and class-wise region embeddings across global images. Each region embedding represents the feature center of one semantic class in an image. We formulate the region embedding of class $c$ by averagely pooling all the pixel embeddings belonging to class $c$ in a single image.

Similar to the Equ. (6), we distill normalized pixel-to-region similarity matrix between the student and teacher networks via KL-divergence loss:

For each mini-batch, we push all teacher region embeddings into the region queue $\mathcal{Q}_{r}$. The overview of our proposed memory-based distillation is shown in Fig. 2.

3 Overall Framework

We summarize our mini-batch-based pixel-to-pixel, memory-based pixel-to-pixel and pixel-to-region distillation together to train the student network. We also employ the conventional pixel-wise cross-entropy task loss $L_{task}$ (Equ. (1)) and class probability KD loss $L_{kd}$ (Equ. (2)) as the basic losses. The overall loss is formulated as:

Here, $\alpha$, $\beta$ and $\gamma$ are weights coefficients. We set $\alpha=1$, $\beta=0.1$ and $\gamma=0.1$. Empirically, we find our CIRKD are not sensitive to coefficients when $\alpha,\beta,\gamma\in[0.1,1]$. When the student and teacher networks mismatch the embedding size, we attach a projection head to the student network. It can map the student’s pixel embeddings to match the teacher’s dimension. The projection head is composed of two $1\times 1$ convolutional layers with ReLU and batch normalization. It would be discarded at the inference phase without introducing extra costs. In Algorithm 1, we use pseudo-code to illustrate the overall training pipeline of CIRKD.

Experiments

Dataset. We employ three popular semantic segmentation datasets to conduct our experiments. (1) Cityscapes is an urban scene parsing dataset that contains 5000 finely annotated images, where 2975/500/1525 images are used for train/val/test. The segmentation performance is reported on 19 classes. (2) CamVid is an automotive dataset that contains 367/101/233 images for train/val/test with 11 semantic classes. (3) Pascal VOC is a visual object segmentation dataset that includes 20 foreground object categories and one background class. We adopt the augmented data with extra annotations provided by . The resulting dataset contains 10582/1449/1456 images for train/val/test.

Evaluation metrics. Following the standard setting, we employ mean Intersection-over-Union (mIoU) to measure the segmentation performance.

Network architectures. For all experiments, we use the segmentation framework DeepLabV3 with ResNet-101 (Res101) backbone as the powerful teacher network. For student networks, we use various segmentation architectures to verify the effectiveness of distillation methods. Specifically, DeepLabV3 and PSPNet with different backbones of ResNet-18 (Res18) and MobileNetV2 (MBV2) are adopted.

Training details. Following the standard data augmentation, we employ random flipping and scaling in the range of $[0.5,2]$. All experiments are optimized by SGD with a momentum of 0.9, a batch size of 16 and an initial learning rate of 0.02. The number of the total training iterations is 40K. The learning rate is decayed by $(1-\frac{iter}{total\_iter})^{0.9}$ following the polynomial annealing policy . For crop size during the training phase, we use $512\times 1024$, $360\times 360$ and $512\times 512$ for Cityscapes, CamVid and Pascal VOC, respectively.

Evaluation details. We evaluate the segmentation performance under a single scale setting over the original image size following the general protocol .

Compared distillation methods. We compare our proposed CIRKD with state-of-the-art segmentation distillation methods: SKD , IFVD and CWD . We re-run all methods using author-provided code. All methods use the same pre-trained teacher DeepLabV3-ResNet101.

Hyper-parameters setup. The hyper-parameters are mainly from the pixel and region queues. For the pixel queue, we set $N_{p}=20K$ for each class and enqueue $V=16$ pixels per class for each image. For the region queue, we set $N_{r}=2K$ for each class. For each mini-batch, we sample $K_{p}=4096$ pixel embeddings from the pixel queue and $K_{r}=1024$ region embeddings from the region queue to compute similarity matrices.

2 Experimental Results

In Table 1, we compare our proposed CIRKD against state-of-the-art distillation methods on Cityscapes in terms of the validation and test mIoU performance. We can observe that all structured KD methods improve student networks under the teacher’s supervision. CIRKD achieves the best segmentation performance across various student networks with similar or different architecture styles. It reveals that CIRKD does not rely on architecture-specific cues. Moreover, our method outperforms the best completing CWD with an average 0.60% validation mIoU gain and 0.78% test mIoU gain across four student networks. The results demonstrate that distilling cross-image relations guides the student to achieve better segmentation performance than intra-image pixel affinity .

As illustrated in Fig. 3, we also show the performance of individual class IoU scores over the student network. We can observe that our CIRKD achieves better class IoU scores than baseline (w/o distillation) and CWD consistently, especially for those categories with low IoU scores. For example, our method obtains 10.4% and 9.4% relative improvements on Wall than baseline and CWD, respectively. We further show the qualitative segmentation results visually in Fig. 4. We can observe that our CIRKD produces more consistent semantic labels with the ground truth than baseline and CWD, indicating more meaningful pixel dependencies are captured.

T-SNE visualization of learned feature embeddings on the student network by CWD and our proposed CIRKD is shown in Fig. 5. Compared to the CWD, the network trained by CIRKD shows a well-structured pixel-wise semantic feature space. The visual result suggests that learning cross-image pixel relations from the teacher network would help the student achieve better intra-class compactness and inter-class separability, thus improving segmentation performance.

2.2 Results on CamVid

In Table 2, we evaluate various distillation methods on CamVid. Our CIRKD achieves the best performance consistently. It outperforms the state-of-the-art CWD by 0.50% and 0.73% mIoU gains over DeepLabV3 and PSPNet, respectively.

2.3 Results on Pascal VOC

Beyond scene-parsing datasets, we also evaluate our CIRKD on Pascal VOC, a representative visual object segmentation dataset. As shown in Table 3, CIRKD achieves the best performance compared to other segmentation KD approaches. It surpasses the best completing CWD by 0.48% and 0.79% mIoU improvements on DeepLabV3 and PSPNet, respectively. The results demonstrate the scalability of our CIRKD to work reasonably well on visual object segmentation.

3 Ablation Study and Parameter Analysis

We conduct thorough ablation experiments of our proposed CIRKD on the Cityscapes validation set, a standard benchmark for semantic segmentation. For all experiments, we choose DeepLabV3-ResNet101 as the teacher and DeepLabV3-MobileNetV2 as the student by default.

Ablation study of loss terms. As shown in Table 4, we examine the contribution of each distillation loss. The conventional KD loss $L_{kd}$ improves the baseline by 1.14%. Applying cross-image relational KD losses of $L_{batch\_p2p}$, $L_{memory\_p2p}$ and $L_{memory\_p2r}$ lead to 0.61%, 0.85% and 0.68% mIoU gains over $L_{kd}$, respectively. The results show two conclusions: (1) Pixel-to-pixel distillation is more informative than the pixel-to-region counterpart. (2) Memory-based pixel-to-pixel distillation is better than the mini-batch-based counterpart, since the former can capture broader pixel dependencies from much more images than the latter. Finally, applying all losses together maximizes the segmentation performance, reducing the gap between the student and teacher from 4.95% to 2.65%.

Impact of the queue size. We investigate the impact of memory sizes of the pixel queue and region queue. As shown in Fig. 6, distillation performance increases as the sizes of the pixel queue and region queue grow. This is because a larger queue could provide more abundant and diverse embeddings for capturing long-range dependencies. The results also show the distillation performance may also saturate at a certain memory capacity.

Impact of the temperature $\tau$. Temperature $\tau$ is used to calibrate the similarity distribution for relational KD. A more significant temperature $\tau$ brings a smoother distribution. As shown in Fig. 7(a), we investigate the impact of $\tau$ in our CIRKD and find $\tau=0.1$ is the best choice.

Impact of the number of contrastive embeddings. As shown in Fig. 7(b) and Fig. 7(c), we examine the number of contrastive embeddings to calculate pixel-to-pixel and pixel-to-region similarity matrices. The distillation performance increases as $K_{p}$ and $K_{r}$ grow, because the similarity distribution with a larger dimension would encode broader pixel dependencies. The upper bound of distillation performance may saturate at $K_{p}=4096$ for pixel-to-pixel distillation and $K_{r}=1024$ for pixel-to-region distillation.

Conclusion

This paper presents a novel cross-image relational KD to transfer global pixel correlations from the teacher to the student for semantic segmentation. Compared to previous KD approaches, our method helps students learn broader pixel dependencies from the teacher. Experiments on public segmentation datasets demonstrate the effectiveness of our CIRKD. We hope our work can inspire future research to explore global pixel relationships for segmentation KD.

Appendix

References