Length-Adaptive Distillation: Customizing Small Language Model for Dynamic Token Pruning (2023)

Abstract: Pre-trained language models greatly improve the performance of various tasks but at a cost of high computation overhead. To facilitate practical applications, there are mainly two lines of research to accelerate model inference: model compression and dynamic computation (e.g., dynamic token pruning). Existing works either adopt these methods individually or simply apply dynamic computation approaches upon a compressed small language model. We argue that they are sub-optimal since the two approaches are separately designed so the compressed model may not be tailored for dynamic computation. To tackle this problem and make compressed small language models faster, we propose Length-Adaptive Distillation, a two-stage knowledge distillation framework that aims to produce a customized small language model for dynamic token pruning. In the general distillation stage, we enforce the student to mimic and reconstruct the teacher’s output based on the dynamically pruned representations. Then in the task-specific distillation stage, the student is further accustomed to token pruning while absorbing the task-specific knowledge. Experimental results on GLUE benchmark demonstrate that our method can make the small language model more customized for dynamic token pruning and achieve better speed-performance trade-off.

## introduction with the rapid progress of pre-trained language models (plms) @xcite @xcite , dramatic improvement has been achieved in a wide range of natural language understanding and generation tasks. despite the remarkable performance those plms have * corresponding author: dongyan zhao. achieved, they generally suffer from high computation overhead, which prevents them from being deployed into resource-scarce applications. hence, there is an urgent need to accelerate inference speed while minimizing performance degradation. to achieve this goal, great efforts have been made, among which model compression and dynamic computation are two major lines of research with different mechanisms. model compression approaches typically convert the large model into a smaller one with reduced layer number and hidden dimension, while the computation graph of the compressed small model is fixed during inference. one of the effective approaches to achieve this goal is knowledge distillation from a large teacher model to a small student model @xcite @xcite . as for dynamic computation methods, they don't change the architecture of the given language model and dynamically prune its computation graph during inference instead. there are two perspectives in dynamic computation for transformer-based @xcite language models: depth-level early exiting @xcite @xcite and width-level token pruning @xcite @xcite . we focus on token pruning methods in this paper. since these two groups of methods achieve inference acceleration from two orthogonal perspectives, a natural idea is to combine them together to obtain faster inference speed. the simplest way to combine them is the pipeline approach @xcite @xcite : preparing a compressed small language model first and applying dynamic computation algorithm upon it next. though proved effective, we argue that the pipeline approach is sub-optimal since model compression and dynamic computation procedures are individually designed and may not be compatible with each other that much. as a result, the potential of such an approach has not been fully exploited. to address this issue and achieve better speed-performance trade-off, we propose a novel knowledge distillation framework named length-adaptive distillation (abbr. lad) that transfers the knowledge from the large language model to a small language model while helping the small model to get adapted to dynamic token pruning. drawing inspirations from @xcite , we adopt a two-stage distillation paradigm where a general-purpose small language model is first distilled on large-scale general corpora and then taskspecific knowledge is injected on task-specific datasets. in general distillation, we mainly consider two issues: (1) how to transfer high-quality general knowledge from the teacher and (2) how to customize the student model for dynamic token pruning. for (1), we propose to transfer tokenlevel knowledge of two types: hidden representations that encode the contextual semantic information of tokens and attention dependency that is generally adopted as the indication of token importance @xcite for token pruning. the alignment of hidden representations between the teacher and the student is achieved by contrastive distillation and the matching of attention dependency is fulfilled by mean square error. while for (2), we employ a teacher with an unpruned computation graph to teach a student with dynamic token pruning where the importance of tokens is measured by the attention scores, and the retention configuration is sampled from a uniform distribution. as only the unpruned tokens have their representations in the student's last layer, we train the student by enforcing the remaining token representations to be close to the teacher's corresponding representations as well as to reconstruct the complete token representations of the teacher's. then in task-specific distillation, we first employ data augmentation to enlarge the task datasets, then distill the student following the dynamic token pruning setting used in general distillation but transfer more task-relevant knowledge by enforcing the sentence embedding of the student to be close to the teacher via contrastive distillation. we conduct experiments on glue benchmark @xcite . we first prove that our knowledge distillation method outperforms ad-vanced knowledge distillation methods on a standard setting (i.e., without token pruning). moreover, we prove that with the help of the customized length-adaptive distillation, our methods successfully take advantage of both model compression and dynamic computation, and achieve dramatically better speed-performance trade-off compared with the pipeline approach. we release our implementation to facilitate future research foot_0 . to sum up, our contributions are three folds: ## methodology our framework adopts a two-stage knowledge distillation paradigm @xcite where a randomly initialized small language model is first distilled with general knowledge, and is then taught with task-specific knowledge. along with the distillation process, the student not only absorbs the knowledge transferred from the teacher, but also gets adapted to dynamic token pruning. after the two distillation stages, we employ evolutionary search @xcite @xcite to obtain a set of length configurations for token pruning with different speedups. in the following subsections, we first introduce our knowledge distillation algorithms in general and task-specific distillation stages, then demonstrate how to apply token pruning upon our model. ## conclusion in this paper, we propose a two-stage knowledge distillation framework lad that transfers general and task-specific knowledge to the student while helping the student to get adapted to dynamic token pruning. we conduct comprehensive experiments on the glue benchmark. the evaluation results prove that our method can effectively take advantage of model compression and dynamic computation and achieve a superior speed-performance trade-off for inference acceleration.