SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiawei Fan*
Intel Labs China
jiawei.fan@intel.com
Chao Li
chao3.li@intel.com
Xiaolong Liu
HoloMatic Technology Co. Ltd.
liuxiaolong@holomatic.com
Meina Song
BUPT
mnsong@bupt.edu.cn
Anbang Yao†
anbang.yao@intel.com
In recent years, knowledge distillation methods based on contrastive learning have achieved promising results on image classification and object detection tasks. However, in this line of research, we note that less attention is paid to semantic segmentation. Existing methods heavily rely on data augmentation and memory buffer, which entail high computational resource demands when applying them to handle semantic segmentation that requires to preserve high-resolution feature maps for making dense pixel-wise predictions. In order to address this problem, we present Augmentation-free Dense Contrastive Knowledge Distillation (Af-DCD), a new contrastive distillation learning paradigm to train compact and accurate deep neural networks for semantic segmentation applications. Af-DCD leverages a masked feature mimicking strategy, and formulates a novel contrastive learning loss via taking advantage of tactful feature partitions across both channel and spatial dimensions, allowing to effectively transfer dense and structured local knowledge learnt by the teacher model to a target student model while maintaining training efficiency. Extensive experiments on five mainstream benchmarks with various teacher-student network pairs demonstrate the effectiveness of our approach. For instance, the DeepLabV3-Res18|DeepLabV3-MBV2 model trained by Af-DCD reaches  $77.03\% / 76.38\%$  mIOU on Cityscapes dataset when choosing DeepLabV3-Res101 as the teacher, setting new performance records. Besides that, Af-DCD achieves an absolute mIOU improvement of  $3.26\% / 3.04\% / 2.75\% / 2.30\% / 1.42\%$  compared with individually trained counterpart on Cityscapes|Pascal VOCi|Camvid|ADE20K|COCO-Stuff-164K. Code is available at https://github.com/OSVAI/Af-DCD.
In computer vision, semantic segmentation is a mainstream dense prediction task aiming to assign the corresponding class to every pixel of input images. Although fundamental models in this task have achieved remarkable progress [1-3], the heavy computational burden and latency of these models severely prohibit their deployment on resource-constrained devices. As a popular solution to this bottleneck, Knowledge Distillation (KD) [4] aims to transfer the knowledge from large models to light-weight ones. Considering KD in supervised semantic segmentation, teachers usually have
*This work was done when the first two authors were interns at Intel Labs China, supervised by Anbang Yao who proposed the original idea and led the project.
Figure 1: The overall framework of Af-DCD, which contains two major parts: (i) masked feature reconstruction; (ii) augmentation-free contrastive distillation loss. Detail illustrations on Af-DCD are shown in Figure 2.
stronger ability of capturing detailed local information of images where the dense and structured knowledge contributes significantly to segment complicated areas, such as boundary and occlusion. Therefore, distillation methods tailored to this task should keep the integrity of dense and structured knowledge, rather than compressing it excessively such as simply using global representations or salient areas. As a result, directly applying traditional KD methods in semantic segmentation, like vanilla KD [4] or feature imitation [5], cannot achieve promising results in most cases.
Thanks to the development of contrastive learning, some distillation methods, like SSKD [6], CRD [7] and G-DetKD [8], employed this paradigm into their designs, which promoted the performance significantly, as they exploited teacher's intrinsic knowledge in the consistencies among different augmentations or various same-category instances. In semantic segmentation, CIRKD [9] designed an implicit contrastive method which aimed to guarantee pixel-pixel and pixel-region consistencies between teacher and student by involving additional contrastive samples. On one side, these existing methods have achieved remarkable progress in image-level or object-level contrastive mimicking. On the other side, however, the efficacy of the aforementioned methods heavily relies on the augmentation and memory buffer, which entail heavy computational and memory cost. Besides, previous contrastive distillation methods have not explicitly modeled relations among pixel-wise or more fine-grained representations within each local patch to transfer teacher's dense and structured knowledge to student, which is critical in semantic segmentation. Thus, when designing a contrastive distillation method for this task, two vital technical problems should be tackled: (i) High resource demands: No matter leveraging augmented samples or storing feature maps in memory buffer, more computational or storage cost is required. For one, the forward of augmented samples incurs extra computational cost. For another, the output feature maps for semantic segmentation are in high resolution. If we want to make dense and structured contrasting, these feature maps should be originally preserved, which occupy a large amount of memories; (ii) Structured knowledge transfer: Although CIRKD defined pixel-wise alignment between student and teacher, it employed contrastive pixel-wise representations from other images, rather than local areas within the same image. It means that no explicit contrasting was defined between student-specific pixel-wise representation to each teacher's pixel-wise representation within local areas, which was agnostic to the structure in teacher's feature. In brief, a contrastive distillation method specially tailored to semantic segmentation, which also entails no extra high resource demands (data augmentation and memory buffer), is essential.
Driven by achieving this target, we first look into the aforementioned two problems and surprisingly discover both of them are incurred from the simple inheritance in the basic definitions of traditional contrastive learning [10, 11]. Therefore, our Augmentation-free Dense Contrastive Knowledge Distillation (Af-DCD) for supervised semantic segmentation tasks can come out, by re-defining these basic concepts: (i) Contrasting samples are not coarse-grained representations of images or objects. Instead, in our Af-DCD, we move further on pixel-level representations and divide each of them into several disjoint partitions (fine-grained representations), which are treated as our contrasting samples; (ii) Positive and negative pairs are not conditioned on categories or objects. Instead, based on the definition of contrasting samples, we define our pairs in a specific teacher-student feature pair  $(F^t,F^s)$ , where fine-grained representations having the same absolute positions in feature maps are used to formulate positive pairs, while representations having different absolute positions but within neighbourhoods are used as negative pairs. Based on these definitions, our Af-DCD can naturally tackle those two problems discussed above. To the first problem, unlike previous methods which had to introduce data augmentation and memory buffer to construct sufficient positive and negative pairs,
our Af-DCD can directly construct these pairs in each local patch without any extra computational and memory cost. To the second problem, we approach it by introducing two new contrastive designs, Spatial Contrasting and Channel Contrasting, which enable student to capture teacher's local dense and structured knowledge by focusing on contextual and positional channel-group information, respectively. As a natural progression, we introduce a hybrid design called Omni-Contrasting, which combines Spatial Contrasting and Channel Contrasting in a neat manner. This design facilitates the simultaneous transfer of both types of information from the teacher to the student, enabling effective augmentation-free contrastive knowledge distillation. Despite leveraging rather dense contrasting, our Af-DCD also performs efficiently. This is due to the patch separation technique we use, which significantly reduces the computational complexity. Additionally, the distance measuring and contrasting calculation can be carried out in parallel, further enhancing the overall efficiency of our method.
Experimental results demonstrate that: (i) Af-DCD exhibits superior performance compared to state-of-the-art methods, on various benchmarks with different teacher-student network pairs; (ii) Af-DCD exhibits even more significant improvements on larger datasets, such as ADE20K, indicating it can enhance student's generalization capability. Moreover, our analytical experiments illustrate that by effectively learning teacher's self-similarity distribution within neighbourhoods and thus reducing the fine-grained feature distances to the teacher, Af-DCD benefits addressing difficult scenarios in semantic segmentation.
Knowledge Distillation. Knowledge distillation can be generally divided into probability-based approach and feature-based approach. Specifically, the former one forces student to mimic teacher's logits as soft-labels [4, 12], while the later one leverages teacher's hidden feature maps or its variants [5, 13, 14] as distillation supervisions. Some recent feature-based approaches [15-17] designed their distillation methods based on masked image modeling mechanism [18, 19] and achieved promising performance in various tasks, as the feature reconstruction can enhance the interdependencies among pixels, which is beneficial to feature distillation.
Knowledge Distillation in Semantic Segmentation. As semantic segmentation is a dense prediction task, the methods specially designed for this task aimed to capture teacher's structured local information. In order to achieve this target, SKD [20] directly measured the similarities of pixel-wise representations between teacher and student, while He et al. [21] leveraged non-local operation with autoencoder to encode local information. CWD [22] evaluated channel-wise pixel distribution contributing to learn teacher's spatial information in each individual channel. Some methods further explored intrinsic knowledge among different samples. For instance, IFVD [23] measured distances from prototypes of different classes and forced student to mimic teacher's intra-class relations, and CIRKD [9] forced student to keep pixel-wise and region-wise consistencies to teacher among various samples in memory buffer.
Contrastive Knowledge Distillation. Inspired by the development of contrastive learning in self-supervised [10, 24] and supervised tasks [11], some recent works leveraged this paradigm and designed new contrastive distillation methods, which made contrasting between teacher's and student's features by employing different views generated from data augmentation [6], or using various samples' features [25, 7] as well as gradients [26] stored in memory buffer. Specifically in dense prediction tasks, G-DetKD [8] constructed ROI feature pairs and executed soft semantic-guided matching which promoted the performance in object detection, and CIRKD [9] designed an implicit contrastive method which leveraged both pixel and region representations to learn structured information in spatial dimension.
Figure 1 illustrates our distillation process in semantic segmentation from a pre-trained teacher  $T$  to a student  $S$  which needs to be trained on a specific dataset. The output features of  $T$  and  $S$  are denoted as  $\hat{F}^t \in \mathbb{R}^{H \times W \times C^t}$  and  $\hat{F}^s \in \mathbb{R}^{H \times W \times C^s}$  respectively, where  $H$  and  $W$  are height and width of
Figure 2: Detailed illustrations on three different types of Af-DCD, which are Spatial Contrasting (top left), Channel Contrasting (top right) and Omni-Contrasting (bottom). For brevity, the contrasting process is illustrated merely using a specific contrastive sample in student feature maps, denoted as  $F_{i,j}^{s}, F_{i,j,k}^{s}, F_{p,i,j,k}^{s}$  in three Af-DCD designs, respectively. The red arrows denote constructing positive pairs, while the blue arrows denote constructing negative pairs. The gray blocks denote other patches which are not considered in calculating the loss in this patch.
feature maps,  $C^t$  and  $C^s$  are the number of channels of teacher and student, respectively. Our method uses a random mask  $M \in \mathbb{R}^{H \times W \times 1}$  to overlap  $\hat{F}^s$  in spatial dimension and then inputs the masked feature into a generator, which can be formulated as  $F^s = f_{\text{generator}}(M \odot \hat{F}^s)$ . Finally, we can get  $F^s \in \mathbb{R}^{H \times W \times C^t}$ , which has the same dimension as teacher's feature. In order to fit general definition, teacher's feature is also processed by a transform  $F^t = f_t(\hat{F}^t)$ . In our design,  $f_t$  is an identity transformation, thus  $F^t = \hat{F}^t$ . After these operations, the reconstructed student feature  $F^s$  has been projected into the same feature space as teacher's feature  $F^t$ . The feature imitation loss  $L_{fd}$  and our Af-DCD loss  $L_{Af - DCD}$  are calculated based on feature pair  $(F^t, F^s)$ . Besides, following the settings in [9, 22], we also employ vanilla logits-based KD loss [4]. In general, the overall loss of our method can be defined as:
$$
L = L _ {t a s k} + \lambda_ {1} L _ {k d} + \lambda_ {2} L _ {f d} + \lambda_ {3} L _ {A f - D C D}, \tag {1}
$$
where  $\lambda_{i}$  is the balancing weight.
3.2 Augmentation-free Dense Contrastive Knowledge Distillation
After illustrating the overall framework, we further explain our core design in detail, which contains two major parts,  $L_{fd}$  and  $L_{Af-DCD}$ . As we mentioned above, both of them are based on the same feature pair  $(F^t, F^s)$ . Therefore, in order to clarify our design, we first answer two key questions: (i) What are the targets of  $L_{fd}$  and  $L_{Af-DCD}$  separately? (ii) How these two losses cooperate with each other? To the first question,  $L_{fd}$  directly help students to imitate teacher's global and salient features, while  $L_{Af-DCD}$  is designed to mimic teacher's dense structured knowledge within local areas, which can further force  $F^s$  to approach  $F^t$  in the micro and fine-grained views. To the second question, in terms of mimicking dense and structured knowledge, these two losses are inextricably bound up with each other. For one,  $L_{fd}$  provides  $L_{Af-DCD}$  an absolute imitation trend, which cannot be implemented in contrastive loss. For another,  $L_{Af-DCD}$  offers  $L_{fd}$  essential constraints on encouraging this process towards mimicking teacher's structured knowledge in each local patch, making student further approach teacher.
Masked Reconstruction based Feature Imitation. The feature imitation loss  $L_{fd}$  is a basic part of our design, which forces student feature  $F^{s}$  to imitate teacher feature  $F^{t}$  directly. Previous methods commonly leverage simple upsampling to align teacher's and student's feature dimensions. Recently, some methods leveraged masked image reconstruction [15, 17], which is a stronger baseline, as it can promote the performance by strengthening the interdependencies among pixels. This process can be formulated as:
$$
F ^ {s} = f _ {\text {g e n e r a t o r}} \left(M \odot \hat {F} ^ {s}\right), \tag {2}
$$
where  $M \in \mathbb{R}^{H \times W \times 1}$  is the mask matrix generated randomly with the mask ratio  $\zeta \in [0,1)$ . Then feature distillation can be formulated as:
$$
L _ {f d} = \sum_ {k = 1} ^ {C ^ {t}} \sum_ {i = 1} ^ {H} \sum_ {j = 1} ^ {W} \left(F _ {i, j, k} ^ {t} - F _ {i, j, k} ^ {s}\right) ^ {2}. \tag {3}
$$
Although masked reconstruction achieves promising performance, one potential weakness may affect its gain in semantic segmentation: Reconstructed features tend to be similar in local areas  ${}^{3}$  . However,in semantic segmentation,local differences indicate structured knowledge,in which student's reconstructed feature maps should keep these differences. Here comes out why we choose masked reconstruction as our basic feature distillation method: (i) choosing a stronger baseline to further test the effectiveness of our Af-DCD; (ii) examining whether our Af-DCD can promote masked reconstruction for generating more structured representations, enhancing feature imitation process.
Motivation and Key Ideas of Af-DCD. In semantic segmentation, we notice two important facts. On one side, different pixels within local areas may contain different semantic information, as they may belong to different categories or different parts of an object. We call this as contextual information. On the other side, differences among channel groups of a pixel representation implicitly indicate the semantic meaning of this pixel, as each output channel is a specific projection of all input feature maps. We call this as positional channel-group information. In short, the contextual and positional channel-group information within each local area in spatial and channel dimensions is of significance in semantic segmentation. If both types of information from teacher's feature maps can be densely utilized as the distillation guidance to facilitate the training of the student model, we conjecture that student's performance and generalization capability can be greatly boosted. However, traditional feature imitation loss  $L_{fd}$  is difficult to capture these two types of information, as they are not directly observable and measurable comparing to salient representation. Aiming to tackle this problem, we define contrastive loss across student's and teacher's pixel-level or more fine-grained representations to explicitly model such kind of knowledge transfer. Specifically, as shown in Figure 2, Af-DCD incorporates three key ideas:
(1) Spatial Contrasting. With the first phenomenon, we leverage pixel-wise representations in  $F^s$  and  $F^t$  and define pixel-wise dense contrasting based on spatial positions, aiming to transfer teacher's spatial contextual information to student;
(2) Channel Contrasting. With the second phenomenon, we then move further on pixel-wise dense contrasting and split every pixel-specific representation into disjoint groups and define group-wise dense contrasting based on channel positions, aiming to force student to learn teacher's positional channel-group information;
(3) Omni-Contrasting. Progressively, we unify the above two contrasting methods in an hybrid design, which takes advantage of tactful feature partitions across both channel and spatial dimensions in local areas by splitting the feature map into disjoint patches of the same size.
**Formulation of Af-DCD.** The formulation of Af-DCD follows the notations in Section 3.1. The projected feature maps of teacher and student are denoted as  $F^{t} \in \mathbb{R}^{H \times W \times C^{t}}$  and  $F^{s} \in \mathbb{R}^{H \times W \times C^{t}}$ , respectively. In the perspective of contrastive learning, we first clarify some important concepts in our Af-DCD:
(1) Samples and Views. Different from existing contrastive distillation methods in classification and detection, the samples in Af-DCD are pixel-level or more fine-grained representations in feature maps of the same image, which are naturally dense. Furthermore, we define that representations of teacher and student are two views which should be aligned by our contrastive loss;
(2) Positive Pairs and Negative Pairs. Under the definitions above, a positive pair can be defined as: a teacher-student sample pair  $(F_{i}^{s}, F_{j}^{t})$  which has the same index  $i = j$ , where  $i$  and  $j$  are general location indices. Similarly, a negative pair can be defined as: a teacher-student sample pair  $(F_{i}^{s}, F_{j}^{t})$  whose indices are different  $i \neq j$ .
<sup>3</sup>Experimental results illustrate that the distances between feature partitions within local areas are usually small (see in Figure 3b).
Based on the above definitions, the loss for a specific sample  $F_{i}^{s} \in F^{s}$  can be formulated as:
$$
l _ {A f - D C D} \left(F _ {i} ^ {s}, F ^ {t}\right) = - \log \frac {\exp \left(- d \left(F _ {i} ^ {s} , F _ {i} ^ {t}\right) / \tau\right)}{\sum_ {j = 1} ^ {N} \mathbb {1} _ {i \neq j} \exp \left(- d \left(F _ {i} ^ {s} , F _ {j} ^ {t}\right) / \tau\right)}, \tag {4}
$$
where  $\tau$  denotes the temperature score,  $N$  is the total number of contrastive pairs and  $\mathbb{1}_{i\neq j}$  is an indicator function which is 1 iff.  $i\neq j$ . Different from traditional contrastive loss that adopts cosine distance on measuring similarities, we use Euclidean distance to measure sample pair similarities, where  $d(F_i^s,F_i^t) = ||F_i^s -F_i^t ||_2^2$ . Our ablative experiments illustrate that improved performance would be attained when we select the same type of the distance function for  $L_{fd}$  and  $L_{Af-DCD}$ .
Spatial Contrasting. As shown in Figure 2 (top left), aiming to learn teacher's spatial contextual information, Spatial Contrasting constructs dense contrasting operations between teacher's and student's pixel-level representations, where for each sample  $F_{i,j}^{s}$ , the positive pair is  $(F_{i,j}^{s}, F_{i,j}^{t})$  while the other  $\{(F_{i,j}^{s}, F_{u,v}^{t})\}_{u,v \neq i,j}$  are all negative pairs. In such condition, we obtain 1 positive pair and  $HW - 1$  negative pairs. Then Spatial Contrasting can be formulated as:
$$
l _ {A f - D C D} ^ {S C} \left(F _ {i, j} ^ {s}, F ^ {t}\right) = - \log \frac {\exp \left(- d \left(F _ {i , j} ^ {s} , F _ {i , j} ^ {t}\right) / \tau\right)}{\sum_ {u = 1} ^ {H} \sum_ {v = 1} ^ {W} \mathbb {1} _ {u , v \neq i , j} \exp \left(- d \left(F _ {i , j} ^ {s} , F _ {u , v} ^ {t}\right) / \tau\right)}. \tag {5}
$$
Channel Contrasting. As shown in Figure 2 (top right), Channel Contrasting first splits each pixel representation into  $M$  non-overlapping channel groups of the same length. Specifically,  $F_{i,j}^{s}$  and  $F_{i,j}^{t}$  are split into  $\{F_{i,j,w}^{s}\}_{w = 1\dots M}$  and  $\{F_{i,j,w}^{t}\}_{w = 1\dots M}$ , where  $w$  is the index of channel group. Then, in order to transfer teacher's channel-group information at each pixel position, the dense contrasting happens on each pixel-level representation pair between  $\{F_{i,j,w}^{s}\}_{w = 1\dots M}$  and  $\{F_{i,j,w}^{t}\}_{w = 1\dots M}$ . For each fine-grained representation  $F_{i,j,k}^{s}$ , the positive pair is  $(F_{i,j,k}^{s}, F_{i,j,k}^{t})$  and negative pairs are  $\{(F_{i,j,k}^{s}, F_{i,j,k}^{t})\}_{w \neq k}$ , where the number of positive and negative pairs is 1 and  $M - 1$ , respectively. Then we substitute these pairs into Formula 4, where  $w$  denotes the index. Channel contrasting can be formulated as:
$$
l _ {A f - D C D} ^ {C C} \left(F _ {i, j, k} ^ {s}, F _ {i, j} ^ {t}\right) = - \log \frac {\exp \left(- d \left(F _ {i , j , k} ^ {s} , F _ {i , j , k} ^ {t}\right) / \tau\right)}{\sum_ {w = 1} ^ {M} \mathbb {1} _ {w \neq k} \exp \left(- d \left(F _ {i , j , k} ^ {s} , F _ {i , j , w} ^ {t}\right) / \tau\right)}. \tag {6}
$$