Prototype-guided cross-task knowledge distillation

Abstract

Recently, large-scale pretrained models have revealed their benefits in various tasks. However, due to the enormous computation complexity and storage demands, it is challenging to apply large-scale models to real scenarios. Existing knowledge distillation methods require mainly the teacher model and the student model to share the same label space, which restricts their application in real scenarios. To alleviate the constraint of different label spaces, we propose a prototype-guided cross-task knowledge distillation (ProC-KD) method to migrate the intrinsic local-level object knowledge of the teacher network to various task scenarios. First, to better learn the generalized knowledge in cross-task scenarios, we present a prototype learning module to learn the invariant intrinsic local representation of objects from the teacher network. Second, for diverse downstream tasks, a task-adaptive feature augmentation module is proposed to enhance the student network features with the learned generalization prototype representations and guide the learning of the student network to improve its generalization ability. Experimental results on various visual tasks demonstrate the effectiveness of our approach for cross-task knowledge distillation scenarios.

1. Introduction

Recently, the Transformer network (Vaswani et al., 2017) has achieved great advances in some visual tasks, for example, image classification (Dosovitskiy et al., 2021; Liu Z et al., 2021; Touvron et al., 2021), object detection (Carion et al., 2020; Zhu XZ et al., 2021; Zhou et al., 2023), image segmentation (Ye LW et al., 2019; Jain et al., 2023), and visual language joint learning (Chen YC et al., 2020; Li LJ et al., 2020; Fu et al., 2023). Based on the self-attention mechanism, Transformer networks can process complete input sequences and possess the advantage of parallelization. Therefore, these networks are usually used to obtain the pretrained model from large-scale datasets (Deng J et al., 2009). Currently, fine-tuning is the common strategy for the utilization of pretrained models in cross-task learning scenarios. After learning the generalized feature representation from large-scale datasets, fine-tuning is performed on the downstream task with the small dataset to boost the performance of the downstream task model. However, applying these large-scale models to practical application scenarios with limited resources (e.g., mobile devices) has become a big challenge due to their enormous computation complexity and huge storage needs.

To solve the above model application issue, some model compression and acceleration technologies have been proposed, for example, parameter pruning (Molchanov et al., 2017; Zhu MH and Gupta, 2018), model quantization (Wu JX et al., 2016), and knowledge distillation (KD) (Hinton et al., 2015). Particularly, KD is a valid approach for model compression, which transfers the knowledge from a large deep neural network into a small one (Hinton et al., 2015). Different from other model compression approaches, KD can decrease the number of parameters and boost the performance of small models on downstream tasks, irrespective of the architectural differences between the teacher network and the student network. It has been successful in a variety of tasks, such as computer vision (Hinton et al., 2015; Romero et al., 2015; Yim et al., 2017; Müller et al., 2019; Park et al., 2019; Gou et al., 2021), natural language processing (Sanh et al., 2019; Sun et al., 2019; Jiao et al., 2020), and speech recognition (Chebotar and Waters, 2016; Kurata and Saon, 2020; Yoon et al., 2021).

However, these KD approaches require mainly the teacher network and the student network to perform the same task; for example, the teacher network and the student network share the same label space, which limits their application in real scenarios, such as downstream tasks in different label spaces as shown in Fig. 1a. By transferring the knowledge from the teacher network to downstream tasks with different label spaces, the cross-task KD method expands the application of the teacher model to a variety of downstream tasks. The existing same-task KD method works mainly to transfer the final prediction logit or the hidden layer knowledge, which is the global-level knowledge alignment and cannot be applied to cross-task KD directly. An earlier work on cross-task KD (Ye HJ et al., 2020) aligns the high-order comparison relationship between models in a local manner; however, this method lags in the representation power of the invariant intrinsic object and is a two-stage distillation method.