Dual-Space Knowledge Distillation for Large Language Models (2024)

Abstract: Knowledge distillation (KD) is known as a promising solution to compress large language models (LLMs) via transferring their knowledge to smaller models. During this process, white-box KD methods usually minimize the distance between the output distributions of the two models so that more knowledge can be transferred. However, in the current white-box KD framework, the output distributions are from the respective output spaces of the two models, using their own prediction heads. We argue that the space discrepancy will lead to low similarity between the teacher model and the student model on both representation and distribution levels. Furthermore, this discrepancy also hinders the KD process between models with different vocabularies, which is common for current LLMs. To address these issues, we propose a dual-space knowledge distillation (DSKD) framework that unifies the output spaces of the two models for KD. On the basis of DSKD, we further develop a cross-model attention mechanism, which can automatically align the representations of the two models with different vocabularies. Thus, our framework is not only compatible with various distance functions for KD (e.g., KL divergence) like the current framework, but also supports KD between any two LLMs regardless of their vocabularies. Experiments on task-agnostic instruction-following benchmarks show that DSKD significantly outperforms the current white-box KD framework with various distance functions, and also surpasses existing KD methods for LLMs with different vocabularies.

## introduction existing large language models (llms) have exhibited strong generalization abilities on various tasks due to their huge model capacities @xcite @xcite . with faith in the scaling law @xcite , the amount of parameters in current llms is expanded steadily to achieve higher intelligence. however, the increasing parameters also bring high deployment costs in real scenarios. for this problem, knowledge distillation (kd; @xcite is one of the promising solutions to compress large models with acceptable performance sacrifice. during the process of kd, the large model typically serves as the teacher and provides supervision signals for a small model (known as the student), and thus the knowledge and the abilities of the teacher can be transferred to the lightweight student. currently, kd algorithms for llms are usually under two frameworks, i.e., black-box kd and white-box kd. black-box kd uses the teacher's decoding sequences as the training data of the student and directly optimizes the cross-entropy loss on the one-hot target. @xcite @xcite . by contrast, white-box kd methods usually minimize the distance (e.g., kl divergence) between the output distributions of the teacher and the student, which theoretically transfer more information and usually perform better than black-box kd @xcite @xcite . although the framework of white-box kd has shown its superiority, the distributions of the student and the teacher in this framework are from different output spaces since they are produced by different prediction heads. at the beginning of this work, we first reveal two inherent limitations in this framework due to the discrepancy of output spaces: on the choice of the distance function d(•||•) in eqn. (2), there have been several explorations (e.g., reverse kl divergence) in recent literature that aim to improve the performance of kd for llms @xcite @xcite . however, in the following section, we will uncover that no matter which distance function is employed, the current white-box kd framework has two inherent limitations since the two distributions p(x i |@xmath0; τ ) and q θ (x i |@xmath1; τ ) are from different output spaces. ## methodology this section introduces our solutions to the above limitations of the current white-box kd framework. firstly, we will introduce our new kd framework in §3.1. then we present a cross-model attention mechanism in §3.2 to extend our framework to support llms with different vocabularies. ## related work white-box kd for language models. the white-box kd framework for language models stems from the standard kd method proposed by @xcite . as pre-trained language models (plms) become prevalent for various nlp tasks, numerous kd methods within this framework were proposed to compress the excessive model sizes of plms @xcite @xcite . besides minimizing the distance between distributions, there are also feature-based kd methods that distill the knowledge in intermediate hidden states and attention maps of the teacher model @xcite @xcite . additionally, white-box kd is also widely used in text generation tasks, such as neural machine translation @xcite @xcite ) and text summarization @xcite . since llms are predominate for various tasks, several kd techniques have also been proposed for llms @xcite @xcite . unlike the previous work that follows the current white-box kd framework, we challenge this framework by revealing its inherent limitations and proposing a simple yet more effective and general kd framework as the solution. kd with the shared prediction head. in the previous literature on kd, simkd @xcite also proposed to share the teacher's prediction head for kd, which was similar to the process of kd in the teacher space in our dskd. however, the aim of simkd is to equip the prediction head of the teacher model to the student model, and thus the student model will be larger after kd and suffer from higher inference costs. in contrast, our dskd only leverages this process to transfer the representation information from the teacher and has no influence on the original model size of the student. ## conclusion in this work, we first reveal two limitations in the current white-box kd framework for llms, i.e., leading to low similarity between the student and the teacher and the requirements of the same vocabulary between two llms. to address them, we propose a novel white-box kd framework, named dual-space knowledge distillation (dskd), which unifies the output spaces of the student and the teacher for kd. on this basis, we further develop a cross-model attention mechanism to solve the vocabulary mismatch between different llms, so that our dskd framework supports kd between any two llms, regardless of their vocabularies. experimental results on several instructionfollowing benchmarks showcase that our framework significantly outperforms the current whitebox kd framework on various distance functions. meanwhile, for llms with different vocabularies, dskd also surpasses all existing kd methods.