Text Representation Distillation via Information Bottleneck Principle (2023)

Abstract: Pre-trained language models (PLMs) have recently shown great success in text representation field. However, the high computational cost and high-dimensional representation of PLMs pose significant challenges for practical applications. To make models more accessible, an effective method is to distill large models into smaller representation models. In order to relieve the issue of performance degradation after distillation, we propose a novel Knowledge Distillation method called IBKD . This approach is motivated by the Information Bottleneck principle and aims to maximize the mutual information between the final representation of the teacher and student model, while simultaneously reducing the mutual information between the student model’s representation and the input data. This enables the student model to preserve important learned information while avoiding unnecessary information, thus reducing the risk of over-fitting. Empirical studies on two main downstream applications of text representation (Semantic Textual Similarity and Dense Retrieval tasks) demonstrate the effectiveness of our proposed approach.

## introduction text representation is a crucial task in natural language processing (nlp) field that aims to map a sentence into a single continuous vector. these representations can be applied to various downstream tasks, such as semantic textual similarity @xcite , information retrieval @xcite ), text classification @xcite , etc. pre-trained language models (plms) have recently become the dominant approach for text representation. however, text representation models developed on large plms typically require enormous computational resources and storage that prevent widespread deployment. further, direct training on small-scale plms will lead to a significant decrease in model performance @xcite . knowledge distillation (kd) is a commonly used approach to reduce the performance gap between large and small models @xcite . this method entails training a large-scale model, referred to as the "teacher" model, followed by transferring its knowledge to a smaller "student" model. the conventional distillation methods mainly concentrate on classification tasks, endeavoring to make the probability distribution of the student model's output as similar as possible to that of the teacher model. recently, various distillation methods have been proposed for representation models @xcite @xcite . the primary objective of these methods is to ensure that the student model's representation closely resembles that of the teacher model. this is typically achieved through the use of a learning objective such as mean squared error (mse) or contrastive learning loss. for an input text x, the final representations of the teacher model and student model are denoted as t and s, respectively. from an information theory perspective, the conventional distillation learning process of continuously approximating t and s through optimization objectives such as mse or contrastive learning can be regarded as maximizing the mutual information (mi) between t and s, which is denoted as i(t, s). however, according to the information bottleneck (ib) principle @xcite , simply maximizing i(t, s) is prone to over-fitting. according to the ib principle, when given an input x and its corresponding label y , our goal is to learn a lowdimensional representation z that is informative for predicting y while minimizing the presence of irrelevant information. the ib method represents this information compressing process as maximizing the mutual information between z and y while minimizing the mutual information between z and x. this approach ensures that the most useful information is preserved, while redundant information is discarded. in text representation distillation, the representations of the teacher and student models, t and s, can be equivalently considered as the label y and the compressed representation z. therefore, we propose a new text representation distillation method based on the ib principle, called ibkd. it aims to maximize the mutual information between s and t while minimizing the mutual information between s and x. as shown in figure 1 , compared to the conventional distillation methods (left), ibkd can effectively filters out task-irrelevant information and retains only the task-relevant information (right), thus improving the generalization of the student model's representation. further, to address the issue of excessive computational effort in directly calculating mutual information @xcite , we introduce different methods to approximate the upper and lower bounds of mutual information. specifically, we use contrastive learning loss @xcite to estimate the lower bound of mutual information and the hilbert-schmidt independence criterion (hsic) @xcite is used to estimate the upper bound. moreover, we found that two-stage distillation can significantly improve the performance of the student model. concretely, the first stage distillation based on large-scale unsupervised data allows the student model to acquire basic text representation characteristics, while the second stage of distillation based on supervised data can further strengthen the representation ability of the student model. importantly, both two stages of distillation can be efficiently performed within the same framework. to demonstrate the effectiveness of our proposed method, we conduct experiments on two main downstream tasks of text representation: the semantic textual similarity (sts) task and the dense retrieval (dr) task. our experimental results have shown that our approach significantly outperforms other methods. briefly, our main contributions are as follows: 1) drawing on the ib principle, we propose a new ibkd method for text representation distillation. 2) we introduce the contrastive learning loss and hsic method to reduce the computational cost of mutual information. 3) we verify the effectiveness of ibkd on multiple benchmark datasets for two different down-streaming tasks. knowledge distillation, initially introduced by @xcite , is a technique employed to convert large, intricate models into smaller, more efficient models while preserving a high level of generalization power. traditional knowledge distillation methods typically utilize a kl divergencebased loss to align the output logits of the "teacher" model and the "student" model @xcite . recently, new approaches have been developed specifically for representation-based models. for instance, the contrastive representation distillation (crd) method proposed by @xcite adopts a contrastive objective to match the representations of the teacher and student models. the hpd method @xcite aims to make the student model's representation similar to a com-pressed representation of the teacher model, thus reducing both the model size and the dimensionality of the final output. the distilcse method @xcite is a two-stage framework that first distills the student model using the teacher model on a large dataset of unlabeled data and subsequently fine-tunes the student model on labeled data. ## method the objective of distillation is to transfer knowledge from a well-trained teacher model (f t ) to a student model (f s ). for input x, the representations of teacher and student are denoted as t and s respectively. referring to the ib principle, we aims to maximize the mutual information between s and t and minimize the mutual information between s and the original input x. formally, the learning objective can be formulated as: however, directly optimizing l ikbd is intractable, especially when x, s, t are high dimensional random variables with infinite support @xcite . consequently, we resort to estimating a lower bound of i(s, t ) and an approximation of i(s, x) instead. in the following subsections, we will first discuss the application of the contrastive learning loss and hsic method in approximating the information bottleneck. subsequently, we will explain how we apply them in our two-stage distillation process. ## conclusion in this paper, we present a new approach for distilling knowledge in text representation called ibkd. our technique is created to address the performance gap between large-scale pre-trained models and smaller ones. drawing inspiration from the information bottleneck principle, ibkd selectively retains crucial information for the student model while discarding irrelevant information. this assists the student model in avoiding over-fitting and achieving a more disentangled representation. through empirical experiments conducted on two text representation tasks, we demonstrate the effectiveness of ibkd in terms of accuracy and efficiency. these results establish ibkd as a promising technique for real-world applications.