Prompting to Distill: Boosting Data-Free Knowledge Distillation via Reinforced Prompt

1 Introduction

Knowledge distillation (KD) has recently emerged as a popular technique for model lightening, finding its applications in a wide spectrum of domains. In some scenarios, however, the training data per se is unavailable, which calls for a variant of KD, known as data-free knowledge distillation (DFKD). Despite the more challenging problem setting, DFKD has also received increasing attention from the community due to its practical nature [Micaelli and Storkey, 2019; Yin et al., 2020; Ma et al., 2020; Fang et al., 2021].

∗Corresponding author

In the literature, [Lopes et al., 2017] initially introduces the idea of data-free knowledge distillation, and follows a learning-by-generating paradigm to craft synthetic data that approximates the original distribution. The key step towards DFKD is to generate high-quality and diverse pseudo samples from a distribution similar to that of the original training set. Once the pseudo samples have been generated, conventional KD algorithms [Sun et al., 2019; Jiao et al., 2020; Hou et al., 2020] can be easily deployed to transfer knowledge from the teacher to the student network.

Prior DFKD algorithms in natural language processing [Ma et al., 2020; Rashid et al., 2021] focused on synthesizing pseudo samples from the teacher's parameters through model inversion [Mahendran and Vedaldi, 2015], where a batch of synthetic utterances or a sentence generator is optimized under the restrictions of some human-crafted distributional priors. The confidence-based prior is the most widely used human-crafted prior for sentence synthesizing. For example, AS-DFD [Ma et al., 2020] aims to find some pseudo samples that can produce high-confidence predictions when fed to the teacher. As shown in Table 1, despite that AS-DFD indeed generates some task-related keywords or phrases that are related to the task, these utterances are still unnatural and of low-quality without correct semantic and syntax.

The above issue is in part caused by the insufficiency of the confidence-based prior in regularizing the synthesized samples. For classification problems, the teacher's predictions mainly contain high-level information like the category of sentences, from which the low-level information (correlation between words, syntax, etc.) is removed. Therefore, the confidence-based prior alone cannot provide robust regularization for the low-level features to make the whole sentences plausible, and leads to some out-of-distribution samples known as "rubbish samples" [Goodfellow et al., 2015]. These samples are far from the underlying distribution of the original training set and are hardly helpful for distillation. Existing methods address this problem by introducing more regularizations into data synthesis. For instance, [Yin et al., 2020] adds extra restrictions using the stored mean and variance statistics in Batch Normalization to improve the quality of synthesized samples. However, these methods are usually inapplicable for language models as they heavily rely on some special architectures like Batch Normalization that are not used in transformer-based models. Therefore, how to introduce more priors to regularize the synthesized sentence is still a challenging problem for DFKD.

To address these problems, we resort to learning natural priors in pre-trained generative language models, and develop a prompting-to-distill framework for DFKD which we term as PromptDFD. Generative language models pre-trained on a large scale of unlabeled corpora grasp the generic prior for languages, which can be incorporated into the synthesis to alleviate the quality issue caused by confidence-based prior. Specifically, the generation of synthesized samples is divided into two steps: 1) a prompt step that explores the language prior in the pre-trained language model in a controllable manner by a task-related topic prompt, and 2) a completion step to expand the prompt into a complete utterance that contains more informative details. To enforce the prompt to be tightly related to the target task when no original data is available, we propose an adversarial reinforced prompt that takes the feedback from both the teacher and the student model into account, simultaneously satisfying the topic and the difficulty requirement. As shown in the experiment, PromptDFD significantly improves the synthesis quality and boosts the student's performance. Extensive experiments demonstrated that the synthesized samples generated by PromptDFD not only contain task-related words, but are meaningful and follow grammatical and linguistic conventions.

In summary, the contribution of this paper can be summarized as follows:

• We propose a novel prompting-to-distill framework for data-free knowledge distillation, through introducing controllable language prior into the synthesized sample.
• We propose a reinforced topic prompter with an adversarial reward. With no access to the original training set, the reinforced prompter can dynamically adjust the topic and the difficulty of the synthesized sample to satisfy the topic constraint of the target tasks
• PromptDFD achieves superior performance on four text classification datasets with different network settings. In some datasets, PromptDFD even yields results on par with those of data-driven KD methods.

2 RelatedWork

Data-driven knowledge distillation. Knowledge distillation (KD), proposed by [Hinton et al., 2015], uses one or more large and cumbersome models (called the teacher model) to train a lightweight model (called the student model), where the student is required to mimic the softened response of the teacher. DistillBERT [Sanh et al., 2019] was the first work to adopt KD on BERT, after which PatientKD [Sun et al., 2019], Tinybert [Jiao et al., 2020], Universal-KD [Wu et al., 2021] and other algorithms have developed other feature alignment criteria, such as intermediate representations, attention maps, extracted relations, etc. Some works [Hou et al., 2020; Chen et al., 2020] focuses on how to dynamically select the most valuable structure of student models during knowledge distillation. Those KD algorithms can be applied as alternative ways to distill knowledge when synthesized samples are obtained.

Data-free knowledge distillation. Existing DFKD algorithms in computer vision adopt the idea of model inversion [Mahendran and Vedaldi, 2015] to reconstruct the training samples from the parameters of the model. Under the constraint of the probabilistic output [Nayak et al., 2019] or historical statistics stored in Batch Normalization [Yin et al., 2020], the training samples [Lopes et al., 2017] or pseudosample generators [Chen et al., 2019] are optimized to explore the underlying distribution of the training data stored in the teacher network. [Micaelli and Storkey, 2019] quantifies the degree of belief matching between the teacher and student. [Fang et al., 2021] proposes to solve the problem of inter-sample pattern collapse using contrast learning. As for methods designed for natural language processing, [Ma et al., 2020] generates pseudo-embeddings to address the discrete problem of tokens, and [Rashid et al., 2021] adopts GumbelSoftmax to pass the gradient to the generator. However, the only constraint that is accessible in BERT is the class prior, which is too weak to reconstruct the underlying distribution.

3 Problem Settings

The goal of data-free knowledge distillation is to reduce redundant parameters in the network without relying on the original dataset. We first start from the data-driven scenario.

Data-driven knowledge distillation. Model compression can be achieved using knowledge distillation, in which the original model is called the teacher network T and the compressed model is called the student network S. Given a sample x = [t1, t2, . . . , tl] ∈ X and its pre-defined label y, where ti is the i-th token in the sample and l is the length of the sentence, the objective of KD can be defined as follows:

LKD = α·LCE(S (x) , y)+(1-α)LKL(T (x) , S (x)) (1) where T (x) and S(x) are the predicted class probabilities on input x. α balances between two criteria, i.e., Cross-Entropy loss LCE and Kullback-Leibler divergence loss LKL.

Data-free knowledge distillation. Data-free knowledge distillation is severely hampered by the lack of access to the initial distribution of data X, and as a consequence, supervisory signals and predefined labels are not available. Thus, constructing the synthesized sample ˆx ∈ ˆ X that can reflect the distribution of X is the key to the data-free knowledge distillation. After the transfer set ˆ X is rebuilt, knowledge distillation techniques can be applied to minimize LKD.

4 Method

4.1 Prompting to Distill

Model inversion with the one-hot class prior fails to constrain the low-level correlation in context, resulting in the semantic confusion of the utterance. The auto-regressive pre-trained language models, such as GPT2, are trained on large-scale corpora to capture these informative co-occurrences. Inspired by [Shin et al., 2020] and [Li and Liang, 2021], our intuition is to query a pre-trained generative language model by prompt engineering and explore the language prior of the language model in a controllable manner.

Step 1: Prompt step. A prefix prompt is used to restrict the topic of the entire synthesized sentence, ensuring that the synthesized sentences are assigned some topics that are relevant to the original training set. Specifically, the topic prompt P1:n = (p1, p2, . . . , pn), where n is the length of the prompts, contains a shortlist of introductory words and can be considered as the controller for the pre-trained generative LM in step 2. We will discuss the design of topic prompts in section 4.2.

Step 2: Completion step. Based on the generated prompt, we adopt a pre-trained auto-regressive generator, named the content generator Gc, to extend the prompt and explicitly display the knowledge embedded in the pre-trained LM. To avoid the knowledge in the pre-trained LM being disturbed when tuning, we freeze the parameter θG of Gc. The goal of Gc is to extend the topic prompt P1:n and generate the content ˆx, which behaves as the synthesized sample in knowledge distillation. Given a half-finished sequence (ˆx1, ˆx2, . . . , ˆxi) generated by Gc, the auto-regressive LM computes the hidden states based on the past tokens and the prompt:

where hi is the hidden state for the temporary token and W projects the hidden state into the distribution space of the tokens. For Gc, we use pre-trained GPT2 [Radford et al., 2019] for content filling. When decoding, a top-k nucleus sampling [Holtzman et al., 2020] is used to encourage the diversity of pseudo samples.

4.2 Prefix Prompt for Topic Controlling

Hand-crafted prompt. We start with a straightforward way of designing a topic prompt by manually selecting some keywords that are pertinent to the task. Table 2 contains a few templates that are designed manually. While hand-crafted prompts are capable of guiding the generation of pseudosamples within a given subject area, there are two major inherent flaws: it is difficult to determine the optimal way to design prompts, and the number of hand-crafted prompts is limited. Though decoding with sampling can provide some variability and randomness in the generation process, the content of the synthesized sample is restricted to limited information.

Reinforced prompt. To solve the problem encountered when designing prompts, we propose to generate prompts automatically. We train a topic prompter Gp, which can be tuned to generate topic prompts that are closely related to the domain of the original dataset.

The generation of prompt P1:n can also be defined as a left-to-right sequence generation problem. Suppose that P1:t = (p1, . . . , pt) has been generated, it's desired to infer the probability distribution of the next token Pr(pt+1|P1:t). Unlike the training objective in the sequence generation problem in general, in DFKD there is no immediate and groundtruth label for Pr(pt+1|P1:t). Furthermore, the teacher model may take a different tokenization method with the content generator Gc, leading to the mismatch in the subwords. Thus, previous gradient-based methods that synthesize embeddings or leverage the Gumbel-Softmax cannot be applied.

We propose a reinforced prompter, which tunes itself through the feedback from both the teacher and the student model. Suppose that the state st for the current moment t is defined as the generated sequence P1:t and the action at as the selection of the next token pt+1. The selecting policy π for the next action is inferred by the prompt generator Gpwith the input being the generated sequence and the output being the probabilities of each token as follows:

ϕG represents the parameter in Gp. It is assumed that the transition is deterministic and that no intermediate reward is given before the prompt is generated. With the discounted factor set to 1, the action-value function QϕG (st, at) equals to the state-value function of the next state VϕG (st+1). To maximize J(ϕG) = VϕG (s1) when QϕG (st, at) = VϕG (st+1), following [Yu et al., 2017; Sutton et al., 1999], the gradient of J(ϕG) can be derived as:

Since an initial token p1 is required for the auto-regressive generator, we choose nine common initial words for p1: The, It, To, There, What, This, All, If and We. In Eq.(4), Q(st, at) is undefined, and we will discuss how to model Q(st, at) to make the topic prompt task-relevant.

Adversarial reward. The teacher and the student model provide delayed estimation for the generated prompts once the synthesized utterance is generated. For state st = P1:t and its corresponding action at = pt+1, the content that expanded based on the temporary topic prompt is as follows:

where ˆx is a complete sentence on which the teacher and the student can estimate its probability for each category. Considering the probability as the likelihood that the synthesized sample belongs to a specific category, the teacher chooses the highest probability among all classes as the returned reward:

where C is the label space. However, unitary feedback from the teacher model makes the prompt generator susceptible to overfitting to a specific prompt, which sometimes corresponds to an out-of-distribution sample (see Figure 2). Thus, an adversarial reward is proposed, where the teacher reward is conjunct with the student reward:

where c′ represents the category with the largest response in the teacher network. Eq.(7) can be considered as an estimation of the discrepancy between the teacher and the student. The discrepancy is narrowed when optimizing the student, while the topic prompter is designed to enlarge this discrepancy. Once the student network has mastered an utterance guided by a certain prompt, the probability of the occurrence for this prompt should be reduced and a lower reward should be assigned. This adversarial reward allows for continuous adjustments of the prompts during the training process, significantly increasing the diversity of the pseudo-samples.

Algorithm 1: The training procedure of PromptDFD

Input: The teacher model T and its parameter θT , the content generator Gc and its parameter θG

Output: The student model S and its parameter θS,

the topic prompter Gp and its parameter ϕG

1 Initial θS with θT 2 for i ← 1 to N do 3 Fix θT , θS, θG 4 Generate P1:n = (p1, p2, . . . , pn) ∼ Gp (·; ϕG) 5 for m ← 1 to n - 1 do 6 Generate ˆx using P1:m+1 by Eq.(5) 7 Compute Q(sm, am) by Eq.(7) 8 end 9 Fix θT and update θS, ϕG 10 Update θS by Eq.(1) 11 Update ϕG by Eq.(4) and Eq.(8) 12 end

Repeat penalty. To avoid the reduction of valid information contained in the prompts due to the repetition of tokens, diversity loss is introduced to maximize the difference between the probability distributions between different tokens:

Algorithm 1 describes the training procedure of PromptDFD.

5 Experiments

5.1 Setup

Datasets. We conduct experiments on four text classification datasets to validate the efficacy of PromptDFD: AG News [Zhang et al., 2015], DBPedia [Auer et al., 2007], IMDb [Maas et al., 2011] and SST-2 [Socher et al., 2013].

Model structure. BERT-base is selected as the structure for the teacher model due to its widespread use. For consistency with previous methods, we use three kinds of structures for the student: BERT6, BERT4, and BERTmini. As for Gc, the pre-trained GPT2 [Radford et al., 2019] is utilized to incorporate informative knowledge into synthesized samples, while a lightweight version of GPT2, DistilGPT2 1 , is selected as Gp to generate prompts that are relatively short in length.

Comparison methods. Considering that the content generator Gc pre-trains on a large scale of unlabeled corpora, which can serve as an additional resource in the data-free setting, we include some baselines incorporating the out-ofdomain data for fair comparison. Baselines are separated into three groups: data-driven knowledge distillation (VanillaKD), data-free knowledge distillation without any external resources (Random Text, Modified-ZSKT, Modified-ZSKD, AS-DFD [Ma et al., 2020]) and data-free knowledge distillation with out-of-domain data (Unlabel-KD and [Rashid et al., 2021]). The introduction of each baseline is given in Appendix .

1 https://huggingface.co/distilgpt2

Implementation. For BERTn, the student uses the first n layers of the teacher network as an initialization scheme. For BERTmini, we follow [Rashid et al., 2021] to take the pretrained BERT-mini to initialize the student. The maximum token length is set to 128. A grid search on parameters is performed in our experiment, where the learning rate of the student is selected from {1e-5, 2e-5, 5e-5, 1e-4}, the learning rate of the Gpis selected from {5e-6, 1e-5, 2e-5}, the batch size is selected from {128, 256}, the epoch for training is set to 10, the temperature of KD is selected from {1, 5, 10}, and α is selected from {0.5, 0.9}. Adam is used to optimize the student network and the topic prompter. The top 50 tokens with the highest probability are selected for decoding and the threshold for top-p is set to 0.95. More implementation details for PromptDFD and Unlabel-KD are given in Appendix.

5.2 Experimental results

We evaluate our proposed methods with two types of prompts: hand-crafted prompts (PromptDFD-Manual) and reinforced prompts (PromptDFD-RL). Table 3 and Table 4 compare the accuracy and the teacher-student agreement of different data-free distillation algorithms. PromptDFD-RL outperforms all other algorithms in four datasets. A significant improvement is observed compared with AS-DFD, largely narrowing the gap between data-driven distillation and data-free distillation. Under the setting that external resources are used, PromptDFD also shows superior performance over [Rashid et al., 2021] and Unlabel-KD. Surprisingly, in SST-2 and DBPedia with BERT6 as student, PromptDFD-RL achieves slightly better performance compared with Vanilla-KD. This indicates that the synthesized set provides sufficient knowledge for students compared with the original training set.

5.3 Quantitative Analysis

Impact of prompt length. Table 5 shows how the length of the prompt influences the distillation performance. The case where the prompt length is extended until it reaches the length of the pseudo sample can be characterized as a direct finetuning of the synthesized sentence. According to the table, the optimal length for prompts is approximately 5-6 words. If the prompt is too short, there will be insufficient prefix constraints in the prompt, while lengthening the prompt results in a decrease in performance due to the increasing complexity of training Gp with a huge action space.

Ablation study. Ablation results are reported in Table 6. Compared with PromptDFD-Manual (designed with templates in table 2), the reinforced prompts can steadily improve the performance, indicating the hand-crafted prompts are not the best choice for pseudo-sample construction. The repeat penalty reduces the redundancy between tokens, showing a 0.12%-0.74% decrease on the performance. Adversarial reward shows a little effect on AG News and IMDb but observes a significant impact on SST-2. The main reason for this is that the AG News and IMDb datasets are relatively simpler than SST-2, which can be demonstrated by the performance of PromptDFD-Manual.

An example of how each module affects the generation of prompts is shown in Figure 2. The original prompt is quite repetitive, which is mitigated by introducing a diversity loss. Nevertheless, we find that the generated prompts are far from the domain of the dataset. The teachers give a high likelihood to OOD samples, but they are not helpful in distilling, which is reflected in the performance of models that are distilled using these samples. While the adversarial reward does not prevent the model from generating such OOD samples, a near-zero reward is given after the student has learned these samples. The topic prompter is biased not to generate those Table 6: Ablation results for PromptDFD. "w/o Adv Reward" denotes the adversarial feedback is replaced by Eq.(6). "w/o Repeat Penalty"

indicates the objective Eq.(8) is not included in the training objective.

prompts, avoiding the prompt generators from getting stuck in out-of-distribution samples.

The quality of the synthesized samples are evaluated by two criteria: the presence of task-related terms and the reasonableness and meaningfulness of the utterance.

Topic relevance of synthesized samples. To visually demonstrate the relationship between the synthesized samples and the original training data, we analyze the frequency of task-related keywords in Figure 3. On the sentiment classification dataset SST-2, PromptDFD tends to generate more emotional words than the unlabeled corpus like wikitext103. These words frequently appear in the original training set, demonstrating that the synthesized set generated by PromptDFD mimics the original data to some extent.

Importance of synthesizing meaningful samples To destroy the semantics in the synthesized samples generated by PromptDFD-RL, we shuffle the words, resulting in a synthesized set that shares the same set of words with broken meanings. Taking these disordered sentences as the distillation set, a noticeable performance drop is observed in Table 7, showing that plausible synthetic utterances are indispensable in data-free knowledge distillation.

6 Conclusion

In this work, we propose a novel data-free knowledge distillation framework PromptDFD that incorporates controllable language prior into the synthesized samples by reinforced prompts. The reinforced prompt is guided by the adversarial feedback to adjust the topic dynamically and increase the diversity. PromptDFD largely improves the quality of the synthesized samples and the distillation performance, and even reaches the performance of data-driven knowledge distillation with the original training data.

Acknowledgments

This work is supported by the Key Research and Development Program of Zhejiang Province, China (No. 2021C01013), the National Key Research and Development Project of China (No. 2018AAA0101900), CKCEST, and MOE Engineering Research Center of Digital Library. This work is also supported by NUS Faculty Research Committee Grant (WBS: A-0009440-00-00) and Advanced Research and Technology Innovation Centre (Project Reference: ECTRP2).