TransKD: Transformer Knowledge Distillation for Efficient Semantic Segmentation

TL;DR

The paper addresses the high computational cost of transformer-based semantic segmentation by introducing TransKD, a transformer-to-transformer knowledge distillation framework that distills both transformer-specific patch embeddings and feature maps from a large teacher to a compact student. It introduces four components—Patch Embedding Alignment (PEA), Cross Selective Fusion (CSF), Global-Local Context Mixer (GL-Mixer), and Embedding Assistant (EA)—and couples them with a unified training objective that combines cross-entropy with patch embedding and feature map losses across four stages, i.e., $L = L_{CE} + \sum_{m=1}^M \alpha_m L_{embd}^m + \sum_{m=1}^M \beta_m L_{fm}^m$ (with $M=4$). TransKD achieves substantial efficiency, reducing FLOPs by over 85% while delivering competitive accuracy, including up to 13.12 percentage points mIoU gains over non-pretrained baselines on Cityscapes and strong performance across Cityscapes, ACDC, NYUv2, and VOC2012. The method demonstrates robustness across architectures and conditions, illustrating a practical path to deploying resource-efficient yet accurate semantic segmentation models for autonomous scene understanding.

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 https://github.com/RuipingL/TransKD.

Figures (15)

• Figure 1: Hard examples for semantic segmentation with knowledge distillation (KD) methods. As compared to the Response-based shu2021channel_distillation and Feature-based chen2021knowledge_review KD methods, our framework, TransKD, enables the model to predict the truck and fence more precisely by exploring the knowledge from both feature maps and patch embeddings.
• Figure 2: Comparison with knowledge distillation frameworks. Our TransKDs compensate for the performance gap between the non-pretrained and pre-trained models effectively while adding negligible parameters. The gains over bars indicate the different amounts of parameters and performances of TransKDs. ImN: pre-training on ImageNet deng2009imagenet. SKD: Structured Knowledge Distillation liu2019structured_knowledge_distillation. KR: Knowledge Review chen2021knowledge_review.
• Figure 3: (a)-(c) Knowledge distillation in computer vision is split into three categories gou2021knowledge_distillation_survey: response-based knowledge distillation, feature-based knowledge distillation, and relation-based knowledge distillation. (d) TransKD extracts the relation-based knowledge of feature maps and transformer-specific patch embedding knowledge at each stage.
• Figure 4: Our knowledge distillation framework TransKD. It is divided into two parts: knowledge distillation of patch embeddings (indicated by red arrows and rectangles) and feature maps (indicated by green arrows and rectangles). The loss function consists of two distillation terms (HCL and MSE) and a cross-entropy term. (a) TransKD-Base is the basic version of TransKD, constructed with two fundamental modules, CSF and PEA. (b) TransKD-GL and (c) TransKD-EA are two optimized versions of TransKD.
• Figure 5: Two-branch architecture of Global-Local Context Mixer (GL-Mixer) in Fig. \ref{['fig:KD_structure']}(b). The global context of an embedding is extracted by multi-head attention, whereas the local features are extracted by convolution operations.