Rethinking Task-Specific Knowledge Distillation: Contextualized Corpus as Better Textbook (2022)

Abstract: Knowledge distillation has been proven effective when customizing small language models for specific tasks. Here, a corpus as ‘textbook’ plays an indispensable role, only through which the teacher can teach the student. Prevailing methods adopt a two-stage distillation paradigm: general distillation first with task-agnostic general corpus and task-specific distillation next with augmented task-specific corpus. We argue that such a paradigm may not be optimal. In general distillation, it’s extravagant to let the diverse but desultory general knowledge overwhelms the limited model capacity of the student. While in task-specific distillation, the task corpus is usually limited and narrow, preventing the student from learning enough knowledge. To mitigate the issues in the two gapped corpora, we present a better textbook for the student to learn: contextualized corpus that contextualizes task corpus with large-scale general corpus through relevance-based text retrieval. Experimental results on GLUE benchmark demonstrate that contextualized corpus is the better textbook compared with jointly using general corpus and augmented task-specific corpus. Surprisingly, it enables task-specific distillation from scratch without general distillation while maintaining comparable performance, making it more flexible to customize the student model with desired model size under various computation constraints.

## introduction pre-trained language models (plms) have achieved remarkable success in a wide range of tasks @xcite @xcite @xcite . however, their satisfactory performance comes at the expense of * corresponding author: dongyan zhao. high computation cost, which makes them not applicable in resource-scarce scenarios. to ease the burden, substantial efforts have been made to compress large plms into small ones with minimum performance degradation, among which we focus on knowledge distillation @xcite . in the literature of language model distillation, a major line of research is the objective functions @xcite @xcite @xcite @xcite @xcite that define how to teach the student. however crucial, there is another factor that plays an indispensable role in knowledge distillation: the corpus as 'textbook', only through which the student can be taught. previous work @xcite proposed a general-then-task-specific distillation paradigm, where the student is first taught on large-scale general corpus by a task-agnostic plm and then distilled on task corpus by an in-task fine-tuned plm. though verified effective, we argue that this paradigm may not be optimal for two reasons. (1) given the downstream task, it's wise to directly transfer the task-specific knowledge instead of overwhelming the student of limited model capacity with diverse but desultory general knowledge derived from general corpus and task-agnostic teacher. (2) when transferring task-specific knowledge, the task corpus is usually limited and narrow, which is not enough for the in-task fine-tuned teacher to transfer its abundant task-specific knowledge before the student begins overfitting. we propose that the silver bullet for customizing a small student model into specific tasks is to transfer as much task-specific knowledge as possible while preventing the student model from overfitting. to achieve this goal, an in-task fine-tuned teacher (aka. task-specific teacher) is indispensable without doubt. but what 'textbook' should we use to teach the student? a large-scale general corpus? abundant and diverse but lacking task-relevance. a task corpus with some kind of data augmentation? highly relevant to tasks but limited and narrow. herein, we propose contextualized corpus, a better 'textbook' that marries up task-relevance with abundance and diversity for task-specific knowledge distillation. the contextualized corpus, as the name implies, is constructed by contextualizing the task corpus with large-scale general corpus through relevance-based text retrieval, enriching the limited task corpus with abundant data that not only holds high task-relevance but also keeps as diverse as general corpus. therefore, it acts as the better textbook with which the in-task fine-tuned teacher can transfer enough task-specific knowledge to the student without worrying about overfitting. we conduct experiments of task-specific knowledge distillation on glue benchmark @xcite . the results demonstrate that the utilization of the proposed contextualized corpus largely improves the performance over the previous method that jointly used general corpus and augmented task corpus. moreover, we find that contextualized corpus enables us to get rid of general distillation as initialization, making it possible to customize small models with desired model size flexibly. our contributions are two folds: (1) we propose contextualized corpus, a better textbook compared with the combination of general corpus and augmentation task corpus, through which task-specific knowledge can be better transferred to student. (2) we conduct comprehensive experiments to study the utilization of contextualized corpus under different distillation settings and provide detailed analyses to demonstrate its superiority. ## experiments datasets. we conduct experiments on glue benchmark @xcite . based on original glue datasets, we construct two types of corpora: taskda and cc. for taskda, we follow the exact data augmentation setting as @xcite . for cc, we first collect the anchor set based on which the relevant sentences are retrieved. for single-sentence tasks (i.e., sst-2, cola), each sentence is an anchor. for sentence-pair tasks (i.e., mnli, qqp, qnli, rte, mrpc, sts-b), we treat each sentence of the pair as a separate anchor. after removing duplicate anchor sentences in the anchor set, we next retrieve top-k relevant sentences from the 160g pre-training corpus of roberta @xcite for each anchor sentence using bm25 @xcite with the extracted keywords as queries @xcite . we choose different top-k for different tasks to make sure cc for each task has similar amounts of data (around 1.6g-1.8g). we build this retrieval system following @xcite . the statistics of various corpus are shown in implementations. for distillation, we adopt roberta large @xcite as the teacher model and a transformer @xcite with 6 layers and 384 hidden dimensions as the student model. when we use cc or wikibook as the textbook to teach the student, the maximum input length is 128, the maximum training step is 400k, the warmup ratio (wr) is 0.01, the weight decay (wd) is 0.01, the batch size (bs) is 256, the learning rate (lr) is 6e-4, following @xcite . as for using taskda, when the student model is randomly initialized, the hyperparameters remain the same for a fair comparison. when the student model has already been distilled (i.e., initialized with general distillation or task-specific distillation on cc), we change the following hyperparameters: lr is 1e-4, wr is 0.06, the maximum training steps for tasks are defined by their maximum training epoch which is 10 for mnli, qnli, qqp, 20 for sst-2, mrpc, rte, sts-b and 50 for cola, following the suggestion by @xcite . when we first fine-tune roberta large and finally fine-tune the task-specific distilled student model, the maximum input length is 128, the training epoch is 10, wr is 0.06 and wd is 0.01. for mnli, qqp, qnli and sst-2, bs is 32, lr is 1e-5. for mrpc, rte, cola and sts-b, we choose lr from {1e-5, 2e-5} and bs from {16, 32}. ## conclusion we study different influential factors of taskspecific knowledge distillation and propose contextualized corpus, a theoretically simple yet highly effective textbook through which the student can better learn task-specific knowledge from teacher. limitations. we improve the performance of task-specific knowledge distillation by proposing contextualized corpus, a better textbook for the student to learn task-specific knowledge from the teacher. though theoretically simple, the construction of contextualized corpus is a bit more complex than traditional data augmentation, which needs a large-scale general corpus as the candidate pool as well as a text retrieval pipeline that should be accurate and efficient.