HRKD : Hierarchical Relational Knowledge Distillation for Cross-domain Language Model Compression (2021)

Abstract: On many natural language processing tasks, large pre-trained language models (PLMs) have shown overwhelming performances compared with traditional neural network methods. Nevertheless, their huge model size and low inference speed have hindered the deployment on resource-limited devices in practice. In this paper, we target to compress PLMs with knowledge distillation, and propose a hierarchical relational knowledge distillation (HRKD) method to capture both hierarchical and domain relational information. Specifically, to enhance the model capability and transferability, we leverage the idea of meta-learning and set up domain-relational graphs to capture the relational information across different domains. And to dynamically select the most representative prototypes for each domain, we propose a hierarchical compare-aggregate mechanism to capture hierarchical relationships. Extensive experiments on public multi-domain datasets demonstrate the superior performance of our HRKD method as well as its strong few-shot learning ability. For reproducibility, we release the code at https://github.com/cheneydon/hrkd .

## introduction large pre-trained language models (plms) (e.g., bert @xcite ) have demonstrated their outperforming performances on a wide range of nlp tasks, such as machine translation @xcite , summarization @xcite , and dialogue generation @xcite . however, their large size and slow inference speed have hindered practical deployments, such as deploying on resourceconstrained devices. to solve the above problem, many compression techniques for plms have been proposed, such as * corresponding author. quantization @xcite , weight pruning @xcite , and knowledge distillation (kd) @xcite . due to the plug-and-play feasibility of kd, it is the most commonly used method in practice, and we focus on it in this work. the purpose of kd is to transfer knowledge from a larger teacher model to a smaller student model @xcite . traditional kd methods only leverage single-domain knowledge, i.e., transferring the knowledge of the teacher model to the student model domain by domain. however, as stated in the purpose of transfer learning, the model performance on target domains can be improved by transferring the knowledge from different but related source domains @xcite , thus the cross-domain knowledge also plays an important role. in addition, several recent works have also proved the advantage of crossdomain knowledge, and many multi-domain kd methods have been proposed. for example, @xcite ; @xcite demonstrate the effectiveness of distilling knowledge from multiple teachers in different domains; @xcite show that jointly distilling the student models of different domains can enhance the performance. nevertheless, these methods fail to capture the relational information across different domains and might have poor generalization ability. to enhance the transferability of the multi-domain kd framework, some researchers have recently adopted the idea of meta-learning. some studies have pointed out that meta-learning can improve the transferability of models between different domains @xcite . for example, meta-kd @xcite introduces an instancespecific domain-expertise weighting technique to distill the knowledge from a meta-teacher trained across multiple domains to the student model. however, the meta-kd framework trains student models in different domains separately, which is inconvenient in real-world applications and might not have enough capability to capture multi-domain correlations. in this paper, we aim to simultaneously capture the relational information across different domains to make our framework more convenient and effective. specifically, we set up several domainrelational graphs to adequately learn the relations of different domains and generate a set of domainrelational ratios to re-weight each domain during the kd process. moreover, since different domains might have different preferences of layer prototypes, motivated by the riesz representation theorem (hartig, 1983), we first construct a set of reference prototypes for each domain, which is calculated by a self-attention mechanism to integrate the information of different domains. then we introduce a hierarchical compare-aggregate mechanism to compare each layer prototype with the corresponding reference prototype and make an aggregation based on their similarities. the aggregated prototypes are finally sent to the corresponding domain-relational graphs. our framework is referred to as hierarchical relational knowledge distillation (hrkd). we evaluate the hrkd framework on two multidomain nlp datasets, including the mnli dataset @xcite and the amazon reviews dataset @xcite . experiments show that our hrkd method can achieve better performance compared with several multi-domain kd methods. we also evaluate our approach under the few-shot learning setting, and it can still achieve better results than the competing baselines. ## method in this section, we detailedly describe the proposed hrkd framework. our hrkd aims to simultaneously capture the relational information across different domains with both hierarchical and domain meta-knowledges. to achieve this goal, we introduce a hierarchical compare-aggregate mechanism to dynamically identify more representative prototypes for each domain, and construct a set of domain-relational graphs to generate re-weighting kd ratios. the overview of hrkd is shown in figure 1 . we first introduce the basic multi-domain kd method in section 2.1, which is a naive framework lacking the ability of capturing cross-domain relations. then we describe the domain-relational graph and compare-aggregate mechanism in section 2.2 and 2.3, respectively, which are the primary modules of our hrkd method to discover the relational information. ## experiment in this section, we conduct extensive experiments on two multi-domain datasets, namely mnli and amazon reviews, to demonstrate the effectiveness of our hrkd method. ## related work pre-trained language model (plm) compression. due to the large size and slow inference speed, plms are hard to be deployed on edge devices for practical usage. to solve this problem, many plm compression methods have been proposed, including quantization @xcite , weight pruning @xcite , and knowledge distillation (kd) @xcite . among them, kd @xcite has been widely adopted due to its plug-and-play feasibility, aiming to distill the knowledge from a larger teacher model to a smaller student model without decreasing too much performance. for example, bert-pkd @xcite distills both intermediate and output layers on fine-tuning. tinybert @xcite additionally distills the embedding layer and attention matrices during pre-training and fine-tuning. meta-kd @xcite proposes to distill knowledge from a crossdomain meta-teacher through an instance-specific domain-expertise weighting technique. in this paper, we propose a novel cross-domain kd framework that captures the relational information across different domains with both domain and hierarchical meta-knowledges, which has a better capability for capturing multi-domain correlations. transfer learning and meta-learning. transfer learning focuses on transferring the knowledge from source domains to boost the model performance on the target domain. among the methods in transfer learning, the shared-private architecture @xcite ) is most commonly applied in nlp tasks, which consists of a shared network to store domain-invariant knowledge and a private network to capture domain-specific information. there are also many works applying adversarial training strategies @xcite @xcite , which introduce domain adversarial classifiers to learn the domain-invariant features. besides, the research of multi-domain learning has gained more and more attention recently, which is a particular case of transfer learning targeting transferring knowledge across different domains to comprehensively enhance the model performance @xcite . unlike transfer learning, the goal of meta-learning is to train a meta-learner that can easily adapt to a new task with a few training data and iterations @xcite . traditional meta-learning typically contains three categories of methods: metric-based @xcite , model-based @xcite , and optimization-based @xcite . in addition, the meta-learning technique can benefit the multi-domain learning task by learning the relationship information among different domains @xcite . in this paper, we leverage meta-learning to solve the multi-domain learning task, where we consider cross-domain kd to simultaneously capture the correlation between different domains, aiming to train a better student meta-learner. ## conclusion in this paper, we present a hierarchical relational knowledge distillation (hrkd) framework to simultaneously capture the cross-domain relational information. we build several domain-relational graphs to capture domain meta-knowledge and introduce a hierarchical compare-aggregate mechanism to capture hierarchical meta-knowledge. the learnt domain-relational ratios are leveraged to measure domain importance during the kd process. extensive experiments on public datasets demonstrate the superior performance and solid few-shot learning ability of our hrkd method.