Plug and Play Knowledge Distillation for k NN - LM with External Logits (2022)

Abstract: Despite the promising evaluation results by knowledge distillation (KD) in natural language understanding (NLU) and sequence-to-sequence (seq2seq) tasks, KD for causal language modeling (LM) remains a challenge. In this paper, we present a novel perspective of knowledge distillation by proposing plug and play knowledge distillation (PP-KD) to improve a (student) kNN-LM that is the state-of-the-art in causal language modeling by leveraging external logits from either a powerful or a heterogeneous (teacher) LM. Unlike conventional logit-based KD where the teacher’s knowledge is built-in during training, PP-KD is plug and play: it stores the teacher’s knowledge (i.e., logits) externally and uses the teacher’s logits of the retrieved k-nearest neighbors during kNN-LM inference at test time. In contrast to marginal perplexity improvement by logit-based KD in conventional neural (causal) LM, PP-KD achieves a significant improvement, enhancing the kNN-LMs in multiple language modeling datasets, showing a novel and promising perspective for causal LM distillation.

## introduction the effectiveness of knowledge distillation (kd) has been extensively validated in natural language processing (nlp) along with various distilled models @xcite @xcite as well as emerging kd approaches @xcite . for causal language modeling, however, it is so rare to see a success of kd as it is in natural language understanding (nlu) and sequence-to-sequence (seq2seq) tasks; even the versatile logit-based kd @xcite , which appears to work in almost any kd scenario with any model architecture with stateof-the-art results @xcite , still does not show a substantial improvement in the metrics of causal language modeling itself (e.g., perplexity) although it may benefit downstream task fine-tuning for a causal lm @xcite . with the motivation to advance the performance boundary, we study the k-nearest neighbor language model (knn-lm) @xcite which is the state-of-art in causal language modeling, and propose plug and play knowledge distillation (pp-kd) to enhance its result, especially for the small-size model, by leveraging knn logits from a teacher lm. unlike conventional logitbased kd where the teacher's knowledge is built-in by training the student with an auxiliary loss to fit the teacher's logits, pp-kd stores the teacher's logits externally and uses them only at test time; thus it is plug and play. as figure 1 shows, pp-kd works during inference to enhance knn results. compared with the vanilla knn-lm, it is required to additionally store the teacher's logits besides context representations and targets of training examples. after retrieving the k nearest neighbors (i.e., training contexts), we get both of their corresponding targets and logits from the datastore and aggregate them into the knn prediction. as pp-kd is plug and play, we can easily enable/disable it by keeping/removing the effects of logits (in the red dashed boxes in figure 1 ) on knn prediction during inference; moreover, we can flexibly switch the teacher we want to employ simply by using its logits instead without retraining like conventional kd. we study pp-kd with two kinds of teachers: one is a more powerful causal lm; the other is a heterogeneous masked lm @xcite . extensive experiments in wiki-103 and bookcorpus demonstrate that pp-kd can significantly benefit causal language modeling, and that a stronger teacher or a teacher ensemble by a causal lm and a masked lm can further improve the perplexity. the contribution of this paper is twofold: training contexts targets logits representations distance steve jobs, ceo of bill gates co-founded steve jobs invented the apple microsoft iphone 3 4 5 testing context targets representations steve jobs is the co-founder of ? classification logits weighting 0.7 0.2 0.1 aggregation interpolation knn retrieval logits lookup target weighting 0.7 0.2 0.1 target lookup figure 1: overview of pp-kd for the knn-lm. the red dashed boxes indicate the external plug and play logits in pp-kd for enhancing knn results at test time, which can be flexibly enabled/disabled or replaced with more informative logits by a more powerful teacher. please note that only logits are from the teacher model; the context representation for knn search is still by the student model. ## conclusion and future work we present pp-kd -a novel perspective for leveraging more powerful (teacher) models to improve state-of-the-art knn-lms for causal language modeling. compared with conventional "built-in" kd, pp-kd leverages the teacher's logits stored externally to enhance the prediction at test time and achieves a rare success in causal lm distillation. as a preliminary and focused study, this work shows promising results of pp-kd in language modeling (as an end task), while it still has much room for improvement (e.g., more efficient implementation, more effective and informative logits as well as more in-depth analyses for pp-kd) and great potential to benefit downstream tasks. we leave these for future work and look forward to building a connection between pp-kd and the emerging retrieval augmented modeling in a bigger picture. borui zhao, quan cui, renjie song, yiyu qiu, and jiajun liang. 2022. decoupled knowledge distillation. arxiv preprint arxiv:2203.08679. yukun zhu, ryan kiros, rich zemel, ruslan salakhutdinov, raquel urtasun, antonio torralba, and sanja fidler. 2015. aligning books and movies: towards story-like visual explanations by watching movies and reading books. in the ieee international conference on computer vision (iccv).