Bridging the Gap between Decision and Logits in Decision-based Knowledge Distillation for Pre-trained Language Models (2023)

Abstract: Conventional knowledge distillation (KD) methods require access to the internal information of teachers, e.g., logits. However, such information may not always be accessible for large pre-trained language models (PLMs). In this work, we focus on decision-based KD for PLMs, where only teacher decisions (i.e., top-1 labels) are accessible. Considering the information gap between logits and decisions, we propose a novel method to estimate logits from the decision distributions. Specifically, decision distributions can be both derived as a function of logits theoretically and estimated with test-time data augmentation empirically. By combining the theoretical and empirical estimations of the decision distributions together, the estimation of logits can be successfully reduced to a simple root-finding problem. Extensive experiments show that our method significantly outperforms strong baselines on both natural language understanding and machine reading comprehension datasets.

## introduction various natural language processing (nlp) tasks have witnessed promising performance from large pre-trained language models (plms) @xcite @xcite . however, plms are usually computationally expensive and memory intensive, hindering their deployment on resource-limited devices. knowledge distillation (kd) @xcite is a popular technique to transfer knowledge from large plms to lightweight models. previous kd works utilize various types of internal information from the teacher model, such as output logits @xcite @xcite , hidden states @xcite , and attention maps @xcite . in real-world applications, however, these types of information are sometimes not accessible due to commercial and privacy issues @xcite . specifically, large-scale plms usually only provide decisions (i.e., top-1 labels) to users. motivated by this scenario, we investigate the task of decision-based kd @xcite for plms, in which only decisions of teacher predictions are available. the information gap between teacher decisions and its internal states is the major challenge for the task. a straightforward approach for decisionbased kd is to treat teacher decisions as ground truth labels and use these labels to train a student model @xcite . however, previous work reveals that logits contain rich knowledge @xcite , relying only on decisions obviously suffers from information loss. to alleviate the problem, @xcite proposes the db3kd method to generate pseudo soft labels according to the sample's robustness. however, db3kd requires that the input of a model can be modified continuously (e.g., image), which hinders its application on plms as their inputs are discrete tokens. therefore, how to fill the information gap under the discrete input setting remains a challenging problem. fortunately, the development of test-time data augmentation for discrete input @xcite @xcite brings hope for resolving the challenge. the basic idea is to modify selected tokens in a piece of text under certain constraints to generate augmented samples and estimate or improve the desired properties of a model based on its behaviors on these samples. test-time data augmentation has been shown to be effective for uncertainty estimation @xcite @xcite , adversary robustness @xcite , and so on. is it possible to narrow down the information gap with test-time data argumentation in decision-based kd for plms? in this work, we propose a novel decision-based kd method for plms. as illustrated in figure 1 , our method is capable of estimating the teacher logits for classes even without observed decisions, narrowing down the information gap between decision and logits. specially, we estimate the logits by combining test-time data argumentation and noncentred orthant probability estimation. on the one hand, we can obtain an empirical estimation of the decision distribution around a sample by test-time data argumentation. on the other hand, we can also derive a theoretical formula for the decision distribution as a non-centred orthant probability, which is a function of logits. as a result, the problem of logits estimation can be reduced to finding the root of the equation that the function takes the value of the empirical estimation. extensive experiments on various natural language understanding and machine reading comprehension datasets demonstrate the effectiveness of our proposed method, which outperforms strong baselines significantly. moreover, quantitative analysis reveals that our method obtains better estimation of logits, narrowing down the information gap. decision-based kd is also related to black-box kd @xcite and distillation-based black-box attacks @xcite @xcite @xcite . both of them involve distilling a student model from black-box models. however, these works generally assume that the score-based outputs are accessible. decision-based kd focuses on a more challenging scenario where only top-1 labels are accessible. ## background knowledge distillation (kd) is a technique that aims to transfer knowledge from the teacher model to the student model by aligning certain statistics, usually the logits, of the student to those of the teacher. given input x, we denote the pre-softmax logits vector of the teacher and student as z and v, respectively. the process of kd involves minimizing the kullback-leibler (kl) divergence between the probabilities induced from z and v as follows: ## conclusion we introduce a novel decision-based kd method, which bridges the information gap between teacher decisions and logits by estimating teacher logits. in contrast to existing solutions for decision-based kd, our method is applicable to nlp tasks with discrete inputs. extensive experiments over various tasks and model architectures demonstrate the effectiveness of our proposed method. one future direction for the decision-based kd is the exploration of other nlp tasks, such as neural machine translation, text generation, and question answering. the other direction is kd from non-nn models to nn models, which benefits the training of nn models with additional information from a wider range of models. unlike conventional kd, decision-based kd does not require internal information from nn models and is promising for solving this problem.