TransKD: Transformer Knowledge Distillation for Efficient Semantic Segmentation

Ruiping Liu,Kailun Yang,Alina Roitberg,Jiaming Zhang,Kunyu Peng,Huayao Liu,Yaonan Wang,Rainer Stiefelhagen
2024-09-05
Abstract:Semantic segmentation benchmarks in the realm of autonomous driving are dominated by large pre-trained transformers, yet their widespread adoption is impeded by substantial computational costs and prolonged training durations. To lift this constraint, we look at efficient semantic segmentation from a perspective of comprehensive knowledge distillation and aim to bridge the gap between multi-source knowledge extractions and transformer-specific patch embeddings. We put forward the Transformer-based Knowledge Distillation (TransKD) framework which learns compact student transformers by distilling both feature maps and patch embeddings of large teacher transformers, bypassing the long pre-training process and reducing the FLOPs by >85.0%. Specifically, we propose two fundamental modules to realize feature map distillation and patch embedding distillation, respectively: (1) Cross Selective Fusion (CSF) enables knowledge transfer between cross-stage features via channel attention and feature map distillation within hierarchical transformers; (2) Patch Embedding Alignment (PEA) performs dimensional transformation within the patchifying process to facilitate the patch embedding distillation. Furthermore, we introduce two optimization modules to enhance the patch embedding distillation from different perspectives: (1) Global-Local Context Mixer (GL-Mixer) extracts both global and local information of a representative embedding; (2) Embedding Assistant (EA) acts as an embedding method to seamlessly bridge teacher and student models with the teacher's number of channels. Experiments on Cityscapes, ACDC, NYUv2, and Pascal VOC2012 datasets show that TransKD outperforms state-of-the-art distillation frameworks and rivals the time-consuming pre-training method. The source code is publicly available at <a class="link-external link-https" href="https://github.com/RuipingL/TransKD" rel="external noopener nofollow">this https URL</a>.
Computer Vision and Pattern Recognition,Robotics,Image and Video Processing
What problem does this paper attempt to address?
Embedding Assistant (EA) and Patch Embedding Alignment (PEA) in TransKD - GL and TransKD - EA. The Embedding Assistant (EA) module is designed to bridge the student and teacher models with the teacher’s number of channels. The Patch Embedding Alignment (PEA) module facilitates the dimensionality transformation of patch embeddings. To address this issue, we introduce the Embedding Assis - tant (EA) module. EA acts as a transformer block with the teacher’s number of channels, allowing the student to learn from a pseudo - teacher - assistant model that combines the student’s transformer blocks and the teacher’s number of channels. The EA module is integrated into the TransKD framework, as shown in Fig. 6, and it is particularly useful when the student model has fewer channels than the teacher model. The EA module helps to build a pseudo - teacher assistant model by combining the student’s transformer blocks and seamlessly bridge teacher and student models. The EA module operates as follows: First, the student’s patch embeddings are aligned with the teacher’s patch embeddings through the PEA module. Then, the aligned student patch embeddings are passed through the EA module to learn from the pseudo - teacher - assistant model. Finally, the processed patch embeddings are used for the next stage of the student transformer. This process ensures that the student can effectively mimic the teacher’s behavior even when there is a significant size difference. C. Feature Map Distillation In addition to