f-Divergence Minimization for Sequence-Level Knowledge Distillation (2023)

Abstract: Knowledge distillation (KD) is the process of transferring knowledge from a large model to a small one. It has gained increasing attention in the natural language processing community, driven by the demands of compressing ever-growing language models. In this work, we propose an FDISTILL framework, which formulates sequence-level knowledge distillation as minimizing a generalized f-divergence function. We propose four distilling variants under our framework and show that existing SeqKD and ENGINE approaches are approximations of our FDISTILL methods. We further derive step-wise decomposition for our FDISTILL, reducing intractable sequence-level divergence to word-level losses that can be computed in a tractable manner. Experiments across four datasets show that our methods outperform existing KD approaches, and that our symmetric distilling losses can better force the student to learn from the teacher distribution.

## introduction increasingly large language models have continued to achieve state-of-the-art performance across various natural language generation tasks, such as data-to-text generation @xcite , summarization @xcite , and dialogue generation @xcite . however, superlarge language models are inaccessible to most users and researchers due to their prohibitively large model size, emphasizing the importance of high-performing, parameter-efficient small neural models. a widely used approach to training small models is knowledge distillation (kd, @xcite , where the small model (known as the student) learns the knowledge from a much larger model (known as the teacher). kd has shown great success in helping smaller models achieve competitive performance across a wide range of applications @xcite @xcite . existing kd approaches can be categorized into two main branches: representation matching and distribution matching. the former aims to imitate the teacher's real-valued intermediate-layer representations, say, with mean squared error @xcite . our work focuses on the latter, distribution matching, where the student model learns the teacher's predictive distribution. @xcite minimize the cross-entropy loss against the teacher-predicted soft labels, which is equivalent to minimizing the kullback-leibler (kl) divergence between the teacher and student. @xcite propose seqkd, arguing that kl divergence should be minimized at the sequence level for language models. however, such an approach tends to learn an overly smooth student distribution to cover the entire support of the teacher distribution due to the asymmetric nature of the kl divergence. this is often known as the mode-averaging problem (figure 1a ). @xcite propose engine, a nonautoregressive translation model that minimizes the energy function defined by the teacher's output distribution. it can be shown that their objective is related to minimizing the reverse kl between the teacher and student (see section 2.2). this, on the other hand, results in the mode-collapsing problem, where the student model is overly concentrated on certain high-probability regions of the teacher distribution (figure 1b ). in this paper, we address knowledge distillation for text generation tasks, and propose f -distill, a unified framework that formulates sequence-level knowledge distillation as minimizing f -divergence functions. existing seqkd @xcite and engine @xcite methods are approximations of kl and reverse kl distillations under the f -distill framework. further, our formulation naturally leads to jensen-shannon (js) divergence and total variation distance (tvd) distillations, where the divergence measures are symmetric in teacher and student distributions. this forces the student to learn the teacher's distribution better, alleviating mode averaging and collapsing problems. we further develop efficient algorithms for our f -distill approach. first, we show that sequencelevel f -divergence can be decomposed step by step either exactly or as an upper bound. second, we propose to sample from the teacher model in an offline manner, mitigating the additional training cost of symmetric divergence measures (namely, js and tvd). we evaluated our approach on four datasets: dart for data-to-text generation @xcite , xsum for summarization @xcite , wmt16 en-ro for machine translation @xcite , and commonsense dialogue @xcite . experiments show that our proposed f -distill variants consistently outperform existing distribution-matching kd methods, allowing f -distill to achieve an add-on performance improvement when combined with representationmatching kd methods. further, results show that our symmetric distilling losses outperform asymmetric ones, confirming that extreme mode averaging or collapsing is not ideal. to sum up, our contributions are three-fold: ## approach in this section, we first review classic knowledge distilling (kd) algorithms and analyze their drawbacks. then, we propose f -distill, a generalized distilling framework for sequence-level distillation. ## related work knowledge distillation (kd) is pioneered by @xcite , who use an ensemble model as the teacher to train a single-model student by minimizing the squared difference between their predicted logits. @xcite propose to directly learn from the output probabilities by minimizing their kl divergence. @xcite propose patient knowledge distillation (pkd), which requires the student to learn from the teacher's intermediate layers. @xcite propose tiny-bert, extending knowledge distillation for transformer models by additional treatments on the attention layers. other recent distilling methods include finding the optimal layer mapping between two models @xcite and learning from multiple teachers @xcite @xcite . the success of kd has since sparked significant interest in its applications to text generation. @xcite investigate sequence-level knowledge distillation (seqkd) for neural machine translation, where they use sampled, hard sequences to approximate the kl divergence. @xcite train a student model by minimizing the energy function defined by a teacher model, which we show is an approximation to reverse kl distillation. @xcite propose imitation-based kd, where the teacher provides oracle probabilities on student-sampled partial sequences to address the exposure bias problem. further, kd has been extensively used to train non-autoregressive text generation models to reduce the complexity of the training data @xcite @xcite . it is noted that our f -distill requires meaningful student sampling and thus is built upon existing kd techniques @xcite , including word-level and intermediate-layer kd. nevertheless, it shows that our approach achieves an add-on performance improvement, and that our contributions are orthogonal to previous work. besides kd, common model compression techniques include parameter pruning and sparse modeling. parameter pruning first trains a dense network and then removes certain neural weights in hopes of not significantly affecting the model performance @xcite @xcite . alternatively, one may apply sparse modeling techniques such as regularization during the training process to ensure zero-valued parameters @xcite @xcite . our work does not follow these directions, as we consider the knowledge distilling setting. regarding the f -divergence function, it has many applications in the machine learning literature. the standard cross-entropy training is equivalent to minimizing the kl divergence between the ground-truth label distribution (often one-hot) and model distribution @xcite . generative adversarial networks @xcite minimize the jensen-shannon divergence by simultaneously training a generator and a discriminator against each other. @xcite minimize α-divergence for adversarial learning, which generalizes kl and rkl, and is a special case of fdivergence functions. @xcite use total variation distance as a regularizer to encourage the model to predict more distinguishable probabilities. further, jsd is used in computer vision kd @xcite , but their tasks do not involve sequential data and the underlying techniques largely differ from our approach. to the best of our knowledge, we are the first to systematically formulate sequence-level knowledge distillation as f -divergence minimization. ## conclusion we propose f -distill, a family of sequence-level distilling methods beyond minimizing the kl divergence. under our framework, we propose and analyze four variants: kl, rkl, js, and tvd distillations, where existing seqkd and engine are approximations of kl and rkl variants; we further derive step-wise decomposition for our f -distill. results on four text generation tasks show f -distill consistently outperforms existing kd methods, and that our symmetric losses (js and tvd) outperform asymmetric ones by avoiding extreme mode averaging and collapsing.