Text Representation Distillation via Information Bottleneck Principle
Authors: Yanzhao Zhang
Dingkun Long
Zehan Li
Pengjun Xie
Venue: EMNLP 2023 Main (2023)
Keywords: knowledge distillation
text representation
text retrieval
language model
TL;DR: A novel Knowledge Distillation method inspired by the Information Bottleneck is proposed to reduce the performance loss during the distillation process of large language models into smaller text representation models.

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 \textbf{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.

Text Representation Distillation via Information Bottleneck Principle
Yanzhao Zhang', Dingkun Long', Zehan Li, Pengjun Xie AAlibaba Group
2yhangyanzhao.2yz, zyz, dingkun. ldk, hatpgjjun.kjj@llllbbbb-iiiii
Abstract
Task-irelevant
Task-irelevant
Pre-trained language models (PLMs) have cently shown great success in text representa- tion field. However, the high computational cost and high-dimensional representation of PLMs pose significant challenges for practi- cal applications.
S-4
Task relevunt
Task relevmnt
Max ((S.))
I -1(5.
make models more ac- cessible, an effective method is distill large models into smaller representation models. In order relieve the issue performance degra- dation after distillation, we propose 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 represen- tation of the teacher and student model, while simultaneously reducing the mutual informa- tion between the student model's 's representation and the input data. This enables the student model to preserve important learned informa- tion while avoiding unnecessary information, thus reducing the risk of over werr-fitinng Empirical studies on two main downstream applications of text representation (Semantic Textual Simi- larity and Dense Retrieval tasks) demonstrate the effectiveness our proposed approach'.
(a) Traditional Method
(b) IBKD
Figure 1: Information diagram teacher embeddings (original circle), student embeddings (dashed red circle) and original input X (green circle). The left side rep- resents the final optimization object the traditional method, while the right side depicts the object of the IBKD method.
widespread deployment. Further, direct training small-scale PL.Ms will lead significant decrease in model performance (Zhao al., 2022).
Knowledge distillation (KD) a commonly used approach to reduce the performance gap between large and small models (Hinton et al., 2015). This method entails training a large-scale model, re- ferred to as the "teacher" model, followed by trans- ferring its knowledge to smaller "student" model. The conventional distillation methods mainly con- centrate classification tasks, endeavoring make the probability distribution of the student model's output similar as possible that of the teacher model. Recently, various distillation meth- ods have been proposed for representation mod- els (Tian et al., 2019; Wu et al., 2021; Zhao et al., 2022). The primary objective of these methods is to ensure that the student model's 's representation closely resembles that of the teacher model. This is typically achieved through the use of learning objective such as mean squared error (MSE) contrastive learning loss.
1 Introduction
Text representation is crucial task in natural lan- guage processing (NLP) field that aims to map a sentence into single continuous vector. These rep- resentations can be applied to various downstream tasks, such as semantic textual similarity (Agirre et al., 2016; Reimers and Gurevych, 2019), infor- mation retrieval (Karpukhin et al., 2020; Long et al., 2022a), text classification (Garg al., 2020), etc. Pre-trained language models (PLMs) have recently become the dominant approach for text represen- tation. However, text representation models devel- oped on large PLMs typically require enormous computational resources and storage that prevent
For an input text. X, the final representations of the teacher model and student model are denoted as and S, respectively. From an information
The source code is publicly available at https://github: com/Al htm/Alibaba-BBe NLte-NLPPIRR NLB/BBD

---

theory perspective, the conventional distillation learning process of continuously approximating T and through optimization objectives such as MSE or contrastive learning can be regarded as maximizing the mutual information (MT) between T and S, which is denoted as (T",S) However, according to the Information Bottleneck (IB) prin- ciple (Tishby and Zaslavsky, 2015), simply maxi- mizing 7 SS) ) is prone to over-fitting According the IB principle, when given an input X and its corresponding label Y, our goal is learn a low- dimensional representation Z that informative for predicting Y while minimizing the presence of irrelevant information. The IB method represents this information compressing process as maximiz- ing 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 informa- tion is discarded. In text representation distillation, the representations of the teacher and student mod- els, and S, can be equivalently considered as the label Y and the compressed representation Z.
model. Importantly, both two stages of distilla- tion 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 re- sults 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.
Related Work
Knowledge Distillation for text representation Pretrained Language Models (PLMs) have demon- strated remarkable success in the field of text repre- sentation (Reimers and Gurevych, 2019; Karpukhin et al., 2020). Recent research has been devoted to enhancing the performance of PLM-based models through specific pretraining tasks (Gao and Callan, 2022: Wu al., 2022; Long et al., 2022b; Shen et al., 2022), contrastive learning (Gao et al., 2021; Karpukhin et al., 2020), and hard negative min- ing (Xiong et al., 2021; Tabassum et al., 2022). However, these methods are primarily designed for large-scale PLMs, and their direct application to small-scale models often leads to a significant performance decline (Zhao al., 2022).
Therefore, we propose a new text representa- tion distillation method based on the IB principle, called IBKD. It aims to maximize the mutual in- formation between S and T while minimizing the mutual information between S and X. As shown in Figure 1, compared to the conventional distil- lation 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 represen-
tation. Further, to address the issue of excessive computational effort in directly calculating mutual information (Alemi et al., 2017a), introduce different methods to approximate the upper and lower bounds of mutual information. Specifically, we use contrastive learning loss (Sordoni et al., 2021) to estimate the lower bound of mutual in- formation and the Hilbert-Schmidt independence criterion (HSIC) (Gretton et al., 2005) is used to estimate the upper bound.
Knowledge distillation, initially introduced by Hinton et al. (2015), is a technique employed to convert large, intricate models into smaller, more efficient models while preserving a high level of generalization power. Traditional knowledge dis- tillation methods typically utilize KL divergence- based loss to align the output logits of the "teacher" model and the "student" model (Zagoruyko and Komodakis, 2017: Sun et al., 2020).
Moreover, we found that two-stage distillation can significantly improve the performance of the student model. Concretely, the first stage distilla- tion based on large-scale unsupervised data allows the student model to acquire basic text represen- tation characteristics, while the second stage of distillation based on supervised data can further strengthen the representation ability of the student
Recently, new approaches have been developed specifically for representation- bbssed models. For instance, the Contrastive Representation Distilla- tion (CRD) method proposed by Tian al. (2019) adopts a contrastive objective to match the repre- sentations of the teacher and student models. The HPD method (Zhao et al., 2022) aims to make the student model's 's representation similar to a com-

---

pressed representation of the teacher model, thus reducing both the model size and the dimensional- ity of the final output. The DistilCSE method (Wu et al., 2021) is a two-stage framework that first dis- tills the student model using the teacher model on a large dataset of unlabeled data and subsequently fine-tunes the student model on labeled data.
However, directly optimizing LIKBD is in- tractable, especially when X, S, T are high dimensional random variables with infinite sup- port (Alemi et al., 2017a). Consequently, we resort to estimating a lower bound of [(S, T) and an ap- proximation 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 our two-stage distillation process.
Information Bottleneck Principle The infor- mation bottleneck (IB) principle (Tishby and Za- slavsky, 2015) refers the tradeoff that exists in the hidden representation between the necessary information required for predicting the output and the information that is retained about the input. has been applied in the study of deep learning dy- namics (Saxe al., 2018; Goldfeld et al.. 2019), resulting in the creation of more interpretable and disentangled representations of data (Bao, 2021; Jeon al., 2021). addition, it has served as a training objective recent works (Belghazi et al., 2018; Paranjape et al., 2020).
3.1 Lower Bound of I(S, T)
To maximize the mutual information T(S, T), we utilize the InfoNCE loss (Liu et al., 2022) as it has been demonstrated to be a lower bound for mutual information (Sordoni et al., 2021). The InfoNCE loss is defined as:
P(S.T) Lnce = log P(S,T) + EsP(S,T)
(2)
However, a significant challenge in the context of Information Bottleneck (IB) is the estimation of the joint distribution of two random variables and the calculation of the entropy of random variable. In response to this challenge, the Variational In- formation Bottleneck (VIB) (Alemi et al., 2017b) 7b) method employs a variation approximation of the original IB, while the ISICC-Botteenek (Ma et al., 2020) method replaces mutual information terms with the Hilbert-Schmidt Independence Criterion (HSIC) (Gretton et al., 2005) to assess the indepen- dence of two random variables. To the best of our knowledge, our research represents the first applica- tion of IB a training objective in the knowledge distillation area.
~ log (1+ M MP(()P(T)
P(S,T)
P(S,T)
I - log((M + +1) + log P(S)
= logg(M + +1) 1)+ + I(S,T)
(3)
where P(S,T) represents the joint distribution of S and T. and P(S) and P(T) represent the marginal distributions of S and 7 respectively. M denotes the number of negative samples. Based on the above equation, it can be inferred that:
[(S () 2 Luce + log M + 1).
(4)
Thus, Ince can be treated as lower bound for I(S, T), and its tightness increases as M grows. In practice, we establish a connection between P(s;, t4) and sim t2) = exp(<?Wt:) where т represents the temperature, W is a learnable matrix used to align the dimensions of 5 and T, and s; and 4 are instances of S and T, respectively. Hence, we have:
3 Method
The objective of distillation is to transfer knowl- edge from a well-trained teacher model (f*)) to a student model (f*). For input X. the representa- tions of teacher and student are denoted as T and S respectively. Referring to the IB principle, we aims 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:
P(S,T)
Lyoe = log P(S,T) EsP(S,T)
(5)
sim 8i, 12)
(1) where I(..)) denotes the mutual information be- tween two random variables and 3 2 is a hyper- parameter controls the tradeoff between two parts.
LIKBD = - I(S, T) + 3x S(
= E,, log
sim 84, t2) Σ3 sim(
1
sim(si, ty) sim 84, 12) + Σn sim(siitt)
log
7
34

---

where denotes the batch size, and we employ other samples within the same batch as negative samples. this particular setting, the number of negative samples M is equal to n.
Input
Ourput
Input
Ousput
Student
S
Mudel
Peor
X
3.2 Approximation of I(
Teacher
Student Model
To minimize the mutual information between X and S. We introduce Hilbert-Schmin Indepen- dence Criterion (HSIC) (Gretton et al., 2005) here as its approximation. HSIC a statistical method used to measure the independence between two random variables:
(a) Distillation Stage
(b) Fine-tuning stage
Figure An illustration of the IBKD training process, with the left depicting the distillation stage and the right depicting the fine-tuning stage, respectively.
ISSI((X,)) = 1|Cxs
performance than directly maximizing mutual in- formation. In our own experiments, we have also observed that the InfoNCE method tends to pro- duce superior results in practice.
where - denotes the Hillert-Sehmidt norm (Gretton et al., 2005), and Cxs is the cross- covariance operators between the Reproducing Ker- nel Hilbert Spaces (RKHSs) (Berlinet and Thomas- Agnan, 2004) of X and S.
3.3 Two Stage Training
Previous studies (Ma et al., 2020) have proven that lower HSIC((, S)) indicates that P(X) and P(S) more independent, and HSIC(X,S)) O if and only if S)
To fully utilize both unsupervised and supervised data, the training process of the student model con- sists of two stages, outlined in prior research (Wu et al., 2021). In the first stage, referred as the distillation stage, we utilize a substantial amount of unlabeled data to train the student model with the following loss function:
0. So minimizing HSIC(X, S) equivalent to minimizing I(X, S). Let D = {111,11)) $1), (x2, .2).................... (zz, $)}} contains samples i.i.d drawn from P(S, X). Gretton et al. (2005) proposed an empirical estimation of HSIC based on D:
LIB - -L,nor B,HSSC((X S)). (8)
HSiC(X, == "
n (6)
during this stage training, the student model can acquire the text representation characteristics of the teacher model, effectively imbuing the former with the latter's features.
the centering matrix is denoted II I 11T
where ir represents the trace matrix. Kx and Ks are kernel Gram matrices (Ham et al., 2004) and respectively, where Kx,, = k<(;, xT) and k(-,-) denotes a kernel function. In this paper, we use the commonly used radial basis function (RBF) kernel (Vert et al., 2004) by experiments
In the second stage, known as the fine-tuning stage, we subject the student model to further fine- tuming using labeled data. This process helps mitigate any potential bias introduced by the use of unlabeled data in the previous stage. Importantly, the labeled data used this stage can be identical to that which was used to train the teacher model.
((;;,)) mr) exp(>1l=;--)) - =z5l),
(7)
where is a hyperparamcter. lighting that the above equation relies solely on positively paired samples to construct the kernel matrices Kx and Ky, which means that the calcu- lation process is notably more efficient than directly estimating mutual information.
is worth high-
Considering labeled dataset consisting of mul- tiple instances, where each instance contains an anchor sample I; a positive sample att and a set of K negative samples: { Zz2..... .... ) In the fine-tuning stage, we continue to utilize combi- nation of the HSIC method and contrastive loss to prevent over-fitting:
Although it is feasible to replace the InfoNCE loss with HSIC to maximize 7 7(5,77 T), previous re- search (Tschannen et al., 2020) has demonstrated that using the InfoNCE loss often leads to better
I, I
(T.sup 32HSTTC(X, (9)

---

where Lice is a supervised InfoNCE loss:
MiniLM-L6 (Wang et al., 2020) as student mod- els following previous works (Wu et al., 2021; Zhao et al., 2022). For the STS task, we employed the state-of- state-of-the-arl thr- model SinCSE-RoBERTature the teacher model, while for the DR task, we uti- lized CoCondenser S as the teacher model. Unless explicitly stated, in the following, the term "Model- IBKD" represents a model that has undergone the distillation training stage, while "Model-IBKD,t refers to the model after the fine-tuning stage's training.
expl 1et Σx exp(
Σ log Ejeap eeep
of o
as
(10)
An overview of the whole training process of IBKD is shown in Figure 2. After the fine-tuning stage, the final student model obtained and can be utilized for downstream tasks.
3.4 Dimension Reduction
Our baseline models include two types: the first involves directly training different-sized pretrained models on the supervised dataset, and the sec- ond involves using previous state-of-the-ee knowl- edge distillation methods for representation models, such the traditional MSE loss based represen- tation distillation method (Kim and Rush, 2016), the HPD (Zhao et al., 2022) method and the CRD method (Tian et al., 2019). In addition, we also fine-tune the HPD model (HPD ft) using the train- ing data and contrastive learning loss for each task based on their public model 6 as a baseline model to estimate the impact of the fine-tuning stage.
In many downstream applications of text represen- tation models, such as information retrieval, the dimension of the text representation model has a di- rect impact on storage costs and search latency. To reduce the final dimension of the student model's representation, during the fine-tuming stage, we add a projection layer after getting S:
S' = W Sr
(11)
where Wp proj E Raxd' d and d' represent the di- mensions of the original student model and the final dimension, respectively. During inference, we will utilize S" as the final representation.
STS
STS,
DR
DR
learning rate
le-4
3e-5
le-4
le-5
4 Experiment
batch size
256
256
128
128
4.1 Experiment Setup
epoch
10
3
3
3
We conduct experiments on two tasks: semantic textual similarity (STS) and dense retrieval (DR). The STS task aims to measure the semantic simi- larity between two sentences, for this task, we eval- uate seven standard datasets: STS12-16 (Agirre al., 2012, 2013, 2014, 2015, 2016), STS-B (Cer al., 2017) and SICK-R (Marelli al., 2014). Following previous work, we use the SentEval toolkit (Conneau and Kiela, 2018) to do the eval- uation and use Spearman's 's rank correlation as the performance metric.
T
0.1
0.05
0.1
0.05
Table I: Hyparameters for STS and DR tasks in different stages. The subscript "ft" indicates the fine-tuning stage, We use the same hyparameters for both TinyBERT-L4 and MiniLM-L6 models.
4.3 Implementation Details
Training Data For the STS task, utilized the same dataset described in Zhao al. (2022) for the first stage training. This dataset comprises the original Natural Language Inference (NLT) dataset along with additional data generated by applying WordNet substitution and back translation to each instance of the NLI dataset. In the fine-tuning stage, we used the "entailment" pairs from the original
The DR task aims to retrieval relevant pas- sages of the given query, for this task we do ex- periments on the MS MARCO Passage Ranking Dataset (Nguyen et al., 2016). We use MRR@ @ 10 and Recall @ 1000 as the evaluation metrics.
4.2 Baseline Models
htttps://uugginggacee co/nreimers/
MiniLM-L6-ttB4--unaae
We select two pretrained models of varying sizes, namely TinyBERT-LA (Jiao et al., 2020) 2 and
"https://hugginggaace
sup -op-smmmmse-rbbertt-laarre
$nttps://hugginggaace co/Luy//
oo-coomenser--nnrro-eeeteees
H-312_v2
'https://hugginnfface co/.uaadeooooo////oooaooaaoa

---

Model
STS12
STS13
STS14
STSIS
STS16
STS-B
SICK-R
Avg
Params
Dimension
SinCSF-RoBERTRbar.
76.53
85.21
80.95
86.03
82.57
85.83
80.50
82.52
110M
768
CAnCSSERRBBRTaaiae Tatarge
77.46
87.27
82.36
86.66
83,93
86.70
81.95
83.76
330M
1024
SimCSF- BERTR
73.02
80.71
76.89
83.01
78.57
81,10
78.19
78.78
14M
312
TinyBERT-MSE
73.07
81.83
77.92
84.49
80.35
81.69
79.10
79.78
14M
312
ThnBBERTHHPD
74.29
83.05
78.80
84.62
81.17
84.36
80,83
81.02
14M
128
TnnBBERT-HPD/
75.17
84.10
79.97
85,44
82.17
85.52
80.65
81.60
14M
128
ThnyBERI-RRD
74.56
83,26
78.71
84.86
80.72
82.11
79.55
80.54
14M
312
TinyBERT BERT-IBKD
74.73
83.56
78.97
84.97
81.68
84.37
80.52
81.69
14M
312
Tiny BhyBBERTIBKD,
76.14
84,45
80,19
85,54
82,51
85,09
80,18
82.01
14M
312
Tiny BERT-IBKD. 128
74.10
83,59
79.38
85.65
81.53
83.87
79.73
81.12
14M
128
SinCSE-MiniLLN
70.34
78.59
75.08
81.10
77.74
79.39
77.85
77.16
23M
384
MiniLM-MSS
73.75
81.42
77.72
83,58
78.99 99
81,19
78.48
79.30
23M
384
MiniLM-HPD
74.94
84.52
80.25
84.87
81.90
84,98
81.15
81.80
23M
128
Minil MiniLM-HPPD M-HPD
76.03
84.71
80.45
85.53
82.07
85.33
80.01
82.05
23M
128
M-CRD
74.79
84.19
78.98
84.70
80.65
82,71
79.91
81,30
23M
384
MiniLM-IBKD
75.57
85.41
80.27
84.99
82.46
84.78
80.48
82.01
23M
384
Minil. M-IBKD
76.77
86.13
81.03
85.66
82.81
86.14
81.25
82.69
23M
384
MiniLM-IBKI--IE
76.34
85.38
81.32
85.34
81.87
85.14
80.67
82.29
23M
128
Table 2: Results on the STS datasets. The teacher model marked with Dimension denotes the output's dimension of the text representation. We bold the best performance of each student's backhone. The results for SimCSE and HPD methods are taken from their respective original papers, while all other results were reproduced by us. TinyBERT BERT-IBKD-12 and MimiLM-IBKD-122 denote the IBKD models fine-tuned with the dimension reduction method, and the final dimension set to 128. The results of IBKD models are statistically significant difference (p < 0.01) compared to other models.
NLI dataset as positive pairs and the "contradiction" pairs negative pairs.
the M-niLM model delivered a Spearman's 's rank correlation performance of 98% while em- ploying just 6.9% of the parameters utilized by SimCSE-RoBERTa agg.ge Remarkably, outper- forms CoancSE-RoEERaatre. which has 4.7 times more parameters. 2) After the distillation training stage, IBKD has already outperformed previous dis- tillation methods, and the subsequent fine-tuning stage yields additional performance gains. For in- stance, fine-tuning led to 0.40% improvement for the TiRyBERT model and 0.68% improvement for the MiniLM student model. 3) Although ap- plying the dimension reduction method reduces performance, it remains competitive performance with previous state-of-hhe-an results.
For the DR task, we utilized all passages and training queries provided by the MS MARCO Pas- sage Dataset during the initial training stage. We then conducted fine-tuning using the labeled data from the same dataset. In this task, the nega- tives were acquired by leveraging the CoCondenser model. More details reported in the Appendix.
Optimizing Setup The values of hyperparame- ters are listed in Table 1. We kept the value of 7 at 0.5, 31 at 1.0, and 32 at 0.5 for experi- ments. All hyperparameters are selected through grid search. The search range for each hyperparam- eters are listed in the Appendix. For the fine-tuning stage, we select hard negatives for each query. We used Adam for optimization. Our code was imple- mented in Python 3.7, using Pytorch 1.8 and Trans- formers 10. All experiments were run on a single 32G NVIDIA V100 GPU. For the DR task, we constructed the index and performed ANN search using the FAISS toolkit (Johnson et al., 2021).
4.5 Experiment Results on DR
Table shows the results on the DR task. We can find that: 1) IBKD effectively reduces the disparity between the student and teacher models compared to previous methods. For example, the MiniLM- IBKD model attains a 97.9% MRR @ 10 perfor- mance with only 4.2% of parameters compared to the teacher model CoCondenser. The smaller pa- rameters make it 4.6 times faster than the teacher model and require only 50% of memory to store the embeddings. 2) The fine-tuning stage has a more significant impact on performance in the DR
4.4 Experiment Results on STS
From the STS results in Table 2, we observe that: 1) Our method outperforms previous meth- ods using the same student model. For instance,

---

Model
MRR @ 10
&ecall@ 1000
Dimension
Params
Speed
Memory
CoCondenser
38.21
98.40
768
110M
500
26G
Tip
28.64
89.97
312
14M
3000
IIG
TinyBERT-MSE
25.98
90.11
312
14M
3000
11G
TinyBERT-CRE
27.40
92.54
312
14M
3000
11G
ThnBERTTHPP
27.77
92.99
128
14M
3000
4.4G
Tiny FhyBEERTHPDD:
34,93
96.04
128
14M
3000
4.4G
TinyBERT-IBKD
28.88
90.02
312
14M
3000
11G
TinyBERT-IBKD)
37.32
97.46
312
14M
3000
IIG
BERT-I--BB3-11
35.57
96.70
128
14M
3000
4.4G
MiniLM-sup
30.51
94.32
384
23M
2300
13G
MiniLM-MSSE
28.12
93.01
384
23M
2300
13G
MiniLM-CRD
28.79
93.12
384
23M
2300
13G
MiniLM-HPD
29.79
93.98
128
23M
2300
4.4G
MiniLM-HPP
36.53
96.70
128
23M
2300
4.4G
MiniLM-IBKD
30.31
93.66
384
23M
2300
13G
MiniLM -IKLMM--B-D-----
37.49
97.81
384
23M
2300
13G
MiniLMMIBBKD-ID
36.32
97.01
128
23M
2300
4.4G
Table 3: Performance of different models on the MS MARCO Passage Ranking Dataset. The teacher model is marked with A. The results of IBKD are statistically significant difference (p < 0.01) compared to other models. We bold the best performance of each student backbone. All results, except for CoCondenser, implementation. The speed denotes the number of sentences encoded by the model per second using one V100.
from our
task. Specifically, when using the MiniLM-L6- HPD model and the MiniLM-IBKD model, the fine-tuning stage resulted in a 5.7%% increase and 7.0% increase MRR 10, respectively. This can be attributed to the fact that the DR task is an asym- metric matching task, and labeled data plays a cru- cial role in guiding the model to accurately derive the semantic association between the query and the document. 3) Applying the dimension reduc- tion method has been observed decrease model performance by approximately 1%. However, this tradcoff is offset by the significant advantage of saving up to 60%-70% of memory costs.
Model TRyBERRTLA-IBKD f
DR
82.01
37.32
TnyBBERT-L-CKD
81.37
35.73
IMQ kernel w linear kernel
81.94
36.83
81.26
36.17
Table Experiment results on the STS and DR task with different kernel method. We report the Average Spearman's rank correlation coefficient of the seven datasets as the performance metric for the STS task and MRR @ 10 for the DR task respectively.
effectiveness of three types of Gram Kernel meth- ods, namely linear, RBF, and inverse multiquadric (IMQ) (Javaran and Khaji, 2012). Our results, pre- sented in Table 4, reveal that the implementation of HSIC loss can significantly enhance the perfor- mance of both STS and DR tasks. Additionally, we observe that the linear kernel leads to decreased performance, while the non-linear kernels (RBF and IMQ) exhibit comparable performance levels.
To ensure our method is robust across teacher models, we conducted experiments for each task using an alternative teacher model. The correspond- ing results are reported the Appendix.
4.6 Ablation Study
Impact of the HSIC loss and Kernel Methods In this study, we aim to evaluate the impact of the HSIC loss and various Gram Kernel methods HSIC. Specifically, we compare the original IBKI) model's performance with a distillation-base ap- proach that only employs the contrastive learning loss in the distillation and fine-tuning stages (re- ferred to as CKD here). Moreover, we examine the
Impact of the 3 value We conducted an analysis to assess the impact of 3 on the performance of IBKD. The results of the STS task are presented in Figure 3, depicting both the distillation stage (in blue) and the fine-tuning stage (in orange).
In the distillation stage, we observed a notable

---

22.0
1.2
$1.6
1.0
HPD181022
$1.4
5.8
OKD181
CAD180.54)
$1.2
$1.0
2.6
1850
02
Simise/78.78
80.8
distiliation staee
20.8
fie-tan stage
5.4
06.1
0.5
1.0
20
8817
-1.6
-15
-13
-12
-10
Figure 3: STS results with different for IBKD the distillation (blue) and fine-tuning (orange) stages.
Figure 5: Alignment and Uniform metric of different Knowledge Distillation method. The darker color de- notes better performance.
Alignment and Uniform We evaluate the qual- ity of embedding of different distillation methods through two widely used metrics: alignment and uniformity Wang and Isola (2020). Alignment mea- sures the expected distance between positive pairs:
10 12
10
(a) CKD
(b) IBKD
Figure 4: The covariance matrix of the student model's representation. The darker color denotes large correla- tion value. Due to space limitation, we only report the results for the first 16 dimensions.
Lalign :- (z,)) Proe(z.,) ( (12)
On the other hand, uniformity measures the ex- pected distances between embeddings of two ran- dom examples:
improvement the performance as 33 increased. However, beyond a certain threshold, the perfor- mance began to decline. Our findings indicate that achieving an optimal balance between maximizing T(S, T) and minimizing 7 (SS, X) can significantly enhance the student model's ability for downstream tasks. Regarding the fine-tuning stage, we noticed that the model's performance remained relatively consistent across different values of B. This obser- vation can likely be attributed to the fact that the model's 's representations have already undergone training to exhibit low mutual information with the input during the preceding distillation stage.
Luniform log ((x,)) Paata(z,y) [e 2(lf(a)-((v))))
(13)
We plot the distribution of the "uniformity- alignment" map for different representation models across different distillation methods in Figure 5.
The uniformity and alignment are calculated on the STS-B dataset. For both uniformity and align- ment, lower values represent better performance. We observe that our IBKD method achieves better trade-off between uniformity and alignment com- pared with other distilled models.
4.7 Analysis
Conclusion
this section, we aim to examine the effect of minimizing mutual information between the stu- dent model and its input. Therefore, we conduct several experiments using the TinyBERT-LA model the student model while maintaining consistent hyperparameters with the primary experiment.
In this paper, we present a new approach for distill- ing knowledge in text representation called IBKD. Our technique created address the performance gap between large-scale pre-trained models and smaller ones. Drawing inspiration from the In- formation Bottleneck principle, IBKD selectively retains crucial information for the student model while discarding irrelevant information. This as- sists the student model in avoiding over-fitting and achieving more disentangled representation. Through empirical experiments conducted text representation tasks, we demonstrate the ef- fectiveness IBKD in terms of accuracy and effi- ciency. These results establish IBKD as a promis- ing technique for real-world applications.
Feature correlation To analyze the feature cor- relation of the CKD and IBKD models, we ran- domly select 10,000 sentences from the Natural Language Inference (NLI) dataset and calculate the covariance matrix for each dimension. As de- picted Figure the representations generated by the IBKD model exhibit low covariance across all dimensions, indicating that IBKD facilitates the learning of a more disentangled representation.

---

6 Limitation and Risk
Computational Semanties, "SEM 2013, June 13-14, 2013, Atlanta, Georgia, USA. pages 32-43. 22-43.Asso0 Associa- tion for Computational Linguistics.
The IBKD method has two primary limitations. Firstly, it requires that the teacher model be a rep- resentation model, which limits the use of other architectures, such as cross-encoder models that take pair of texts as input and output their se- mantic similarity score. Secondly, the fine-tuning stage need labeled data which may not be avalia- bel in certain situation, even though these data can be the same data as that used to train the teacher model. Additionally, our model may perpetuate biases, lack transparency, pose security and privacy risks, similar to other pre- trained models.
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. 2017a. Deep variational informa- tion bottleneck. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Pro- ceedings. OpenReview.net.
Alexander A. Alemi, lan Fischer, Joshua V. Dillon, and Kevin Murphy. 201 Deep variational informa- tion bottleneck. In Sth International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Pro- ceedings. OpenReview.net.
Feng Bao. 2021. Disentangled variational information bottleneck for multiview representation learning. In Artificial Intelligence - First CAAI International Con- ference, CICAI 2021, Hangzhou, China, June 5-6, 2021, Proceedings, Part II, volume 13070 of Lecture Notes in Computer Science, pages 91-102. Springer.
References
Eneko Agirre, Carmen Banea, Claire Cardie, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agirre, Weiwei Guo, fñigo Lopez-Gazpio, Montse Maritxalar, Rada Mihalcea, German Rigau, Larraitz Uria, and Janyce Wiebe. 2015. Semeval-2015 task 2: Semantic textual similarity, english, spanish and pilot on interpretahil- ity. In Proceedings of the 9th International Work- shop on Semantic Evaluation, SemEval @ NNACL- HLT 2015, Denver, Colorado, USA, June 4-5, 2015, pages 252-263. The Association for Computer Lin- guistics.
Ishmael Belghazi, Sai Rajeswar, Aristide Baratin, R. Devon Hjelm, and Aaron C. Courville. 2018. MINE: mutual information neural estimation. CoRR. abs/1801. 04062..
Alain Berlinet and Christine Thomas-Agnan, 2004. Re- producing kernel hilbert spaces probability and statistics.
Eneko Agire, Carmen Banea, Claire Cardie, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agine Weiwei Guo, Rada Mihalcea, German Rigau, and Janyce Wiebe. 2014. Semeval-2014 task 10: Multilingual semantic textual similarity. In Proceedings of the 8th International Workshop on Semantic Evaluation. SemFval@ COLING 2014, Dublin, Ireland, August 23.24, 2014, pages 81-91. The Association for Com- puter Linguistics.
Daniel M. Cer, Mona T. Diab, Eneko Agirre, Inigo and Lucia Specia. 2017. Semeval- 2017 task I: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 11th International Workshop on Semantic Eval- uation, SemEval ACL 2017, Vancouver, Canada, August 3-4, 2017, pages 1-14. Association for Com- putational Linguistics.
Eneko Agirre, Carmen Banea, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agirre, Rada Mihalcea, Ger- man Rigau, and Janyce Wiebe. 2016. Semeval-2016 task 1: Semantic textual similarity. monolingual and cross- loosulli evaluation. In Proceedings of the 10th International Workshop on Semantic Evaluation, CemEval@NNACL-HLL NAACL-HLT 2016, San Diego, CA, USA, June 16- 17, 2016, pages 497-511. The Association For Computer Linguistics.
Alexis Conneau and Douwe Kiela. 2018. Senteval: An evaluation toolkit for universal sentence representa- tions. In Proceedings of the Eleventh International Conference on Language Resources and Evaluation, LREC 2018, Miyazaki, Japan, May 7-12, 2018. Euro- pean Language Resources Association (ELRA).
Luyu Gao and Jamic Callan. 2022. Unsupervised cor- pus aware language model pre-training for dense pas- sage retrieval. In Proceedings of the 60th Annual Meeting of the Association for Computational Lin- guisties (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pages 2843-28533 Associ- ation for Computational Linguistics.
Eneko Agirre, Daniel M. Cer, Mona T. Diab, and Aitor Gonzalez-Agirre. 2012. Semeval-2012 task 6: A pilot on semantic textual similarity. In Proceed- ings of the 6th International Workshop on Semantic Evaluation, SemEvul @ NAACL-HLT 2012, Montréal, Canada, June 7-8, 2012, pages 385-393. The Asso- ciation for Computer Linguistics.
Tianyu Gao, Xingcheng Yao, and Danqi Chen. 2021. Simese: Simple contrastive learning of sentence em- beddings. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Domini- can Republic, 7-11 November, 2021, pages 6894 6910. Association for Computational Linguistics.
Eneko Agirre, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agire, and Weiwei Guo. 2013, *sem 2013 shared task: Semantic textual similarity. In Proceed- ings of the Second Joint Conference on Lexical and

---

Siddhant Garg, Rohit Kumar Sharma, and Yingyu Liang. 2020. Simpletran: Transferring pre-trained sen- tence embeddings for low resource text classification. CoRR, abs/2004.051199
and Wen-tau Yih. 2020. Dense passage retrieval for open- open-domain question answering. In Proceedings of the 2020 Conference on Empirical Methods in Nat- ural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 6769-6781 -.sscii Associa- tion for Computational Linguistics.
Ziv Goldfeld, Ewout van den Berg, Kristjan H. Gree- newald, Igor Melnyk, Nam Nguyen, Brian Kingsbury, and Yury Polyanskiy. 2019. Estimating information flow in deep neural networks. In Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California. USA, volume 97 of Proceedings of Machine Learning Research, pages 2299-2308 PMLR.
Yoon Kim and Alexander M. Rush. 2016. Sequence- level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pages 1317-1327. The Association for Computational Linguistics.
Arthur Gretton, Olivier Bousquet, Alexander J. Smola, and Bernhard Schölkopf. 2005. Measuring statistical dependence with hilhert-schmidt norms. In Algorith- mic Learning Theory, 16th International Conference, ALT 2005, Singapore, October 8-11, 2005, Proceed- ings, volume 3734 of Lecture Notes in Computer Science, pages 63-77. Springer.
Yue Liu, Junqi Ma. Yufci Xie, Xuefeng Yang, Xingzhen Tao, Lin Peng, and Wei Gao. 2022. Contrastive pre- dictive coding with transformer for video representa- tion learning. Neurocomputing, 482: 152-
Dingkun Long, Qiong Gao, Kuan Zou, Guangwei Xu, Pengjun Xie, Ruijic Guo, Jian Xu, Guanjun Jiang, Luxi Xing. and Ping Yang. 2022a. Multi-cpr: A multi domain chinese dataset for passage retrieval. I In SIGIR '22: The 45th International ACM SIGIR Conference on Research and Development in Infor- mation Retrieval, Madrid, Spain, July 11 15, 2022, pages 3046-3056. ACM.
Jihum Ham. Daniel Lee, Sebastian Mika, and Bern- hard Schölkopf. 2004. A kernel view of the dimen- sionality reduction of manifolds. In Machine Learn- ing. Proceedings of the Twenty-first International Conference (ICML 2004), Banff, Alberta, Canada, July 4-8, 2004, volume 69 of ACM International Con- ference Proceeding Series. ACM.
Dingkun Long, Yanzhao Zhang, Guangwei Xu, and Pengjun Xie. 2022b. Retrieval oriented masking pre- training language model for dense passage retrieval. CoRR, abs/2210. 15133.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge in neural network. CoRR, abs//15330025.
Kurt Wan-Duo Ma. J. P. Lewis, and W. Bastiaan Kleijn. 2020. The HSIC bottleneck: Deep learning without back-propagation. In The Thirty-Fourth AAAI Con- ference Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial In- telligence, EAAI 2020, New York, NY. USA, February 7-12, 2020, pages 5085-5092. AAAI Press.
Saleh Hamzehei Javaran and Naser Khaji. 2012. Inverse multiquadric ( imq) ) function as radial basis function for plane dynamic analysis using dual reciprocity boundary element method.
Insu Jeon, Wonkwang Lee, Myeongjang Pyeon, and Gunhee Kim. 2021. IB IB-GAN: disentangled repre- sentation learning with information bottleneck gen- erative adversarial networks. In Thirty-Fifth AAAT Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial In- telligence, EAAI 2021, Virtual Event, February 2-9, 2021, pages 7926-79344 1934.99..444 AAAI Press.
Marco Marelli, Stefano Menini, Marco Baroni, Luisa Bentivogli, Raffaella Bernardi, and Roberto Zam- parelli. 2014. A SICK cure for the evaluation of compositional distributional semantic models. In Proceedings of the Ninth International Conference on Language Resources and Evaluation. LREC 2014, Reykjavik, Iceland, May 26-31, 2014, pages 216-223. European Language Resources Association (ELRA).
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020. Tinybert: Distilling BERT for natural language un- derstanding. In Findings of the Association for Com- putational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020, volume EMNLP 2020 of Find- ings of ACL, pages 4163 4174. Association for Com- putational Linguistics.
Tri Nguyen, Mir Rosenberg, Xia Song. Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. 2016. MS MARCO: A human generated machine reading comprehension dataset. In Proceedings of the Workshop on Cognitive Computation: Integrat- ing neural and symbolic approaches 2016 co-located with the 30th Annual Conference on Neural Infor- mation Processing Systems (NIPS 2016), Barcelona, Spain, December 9, 2016. volume 1773 of CEUR Workshop Proceedings. CEUR-WS.org
Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2021. Billion-scale similarity search with gpus. JEEE Trans, Big Data, 733)5535-547.
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick S. H. Lewis, Ledell Wu, Sergey Edunov, Danqi Chen,
Bhargavi Paranjape, Mandar Joshi, John Thickstun, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2020.

---

An information bottleneck approach for control- ling conciseness in rationale extraction. CoRR, abs/2005.00652.
Tongzhou Wang and Phillip Isola. 2020. Understanding contrastive representation learning through alignment and uniformity the hypersphere. In Proceedings of the 37th International Conference on Machine Learn- ing, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 9929-9939. PMI..
Nils Reimers and Iryna Gurevych. 2019. Sentence- Sentence embeddings using siamese bett-networks. In Proceedings of the 2019 Conference on Empiri- cal Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing. EMNNLP-UCNIPP 2019, Hong Kong, China, November 3-7, 2019, pages 3980-3990. Association for Computational Linguistics.
Wenhui Wang. Furu Wei, Li Dong, Hangbo Bao, Nan Yang. and Ming Zhou. 2020. Minilm: Deep self- attention distillation for task-agnostic compression of pre pre-trained transformers. In Advances in Neural Information Processing Systems 33: Annual Confer- ence Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
Andrew M. Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan D. Tracey, and David D. Cox. 2018. On the information bottleneck theory deep learning. In 6th International Con- ference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 2018, Conference Track Proceedings. OpenReview.net
Xing Wu, Chaochen Gao. Jue Wang. Liangjun Zang- Zhongyuan Wang, and Songlin Hu. 2021. Disco: Ef- fective knowledge distillation for contrastive learning of sentence embecldings. CoRR, abs/2112.056388.
Tao Shen, Xiubo Geng. Chongyang Tao, Can Xu, Xiao- long Huang, Binxing Jiao, Linjun Yang. and Daxin Jiang. 2022. Lexmae: Lexicon- bottleeekkdd pretrain- ing for large-scale retrieval. CoRR, abs/2208. 14754.
Xing Wu, Guangyuan Ma. Meng Lin, Zijia Lin. Zhongyuan Wang, and Songlin Hu. 2022. Contex- tual mask auto-encoder for dense passage retrieval. CoRR, abs/2208 076700.
Alessandro Sordoni, Nouha Dziri, Hannes Schulz, Ge- offrey J. Gordon, Philip Bachman, and Remi Tachet des Combes. 2021. Decomposed mutual information estimation for contrastive representation learning. In Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Vir- tual Event, volume 139 of Proceedings of Machine Learning Research, pages 9859-9869. PMLR.
Shitao Xiao, Zheng Liu, Yingxia Shao, and Zhao Cao. 2022. Retromae: Pre- Pre-training retrieval- -oriened lan- guage models via masked auto-encoder. In Proceed- ings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pages 538-548. Association for Computational Lin- guistics.
Siqi Sun, Zhe Gan, Yuwei Fang. Yu Cheng, Shuohang Wang, and Jingjing Liu. 2020. Contrastive distil- lation on intermediate representations for language model compression. In Proceedings of the 2020 Con- ference Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 498-508. Association for Computational Linguistics.
Lee Xiong, Chenyan Xiong, Ye Li, Kwok- Tang. Jialin Liu, Paul N. Bennett, Junaid Ahmed, and Arnold Overwijk. 2021. Approximate nearest neigh- bor negative contrastive learning for dense text re- trieval. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Afrina Tabassum, Muntasir Wahed, Hoda Eldardiry, and Ismini Lourentzou. 2022. Hard negative sam- pling strategies for contrastive representation learn- ing. CORR, abs/2206.01197.
Sergey Zagoruyko and Nikos Komodakis. 2017. Paying more attention attention: Improving the perfor- mance of convolutional neural networks via attention transfer. In 5th International Conference on Learn- ing Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. 2019. Contrastive representation distillation. CoRR, abs/1910.10699
Xuandong Zhao, Zhiguo Yu, Ming Wu, and Lei Li. 2022. Compressing sentence representation for semantic re- trieval via homomorphic projective distillation. Findings of the Association for Computational Lin- guisties: ACL 2022, Dublin, Ireland, May 22-27, 2022, pages 774-781, Association for Computational Linguistics.
Naftali Tishby and Noga Zaslavsky. 2015. Deep learn- ing and the information bottleneck principle. CoRR, abs/15033 500./500..2406.
Michael Tschannen, Josip Djolonga, Paul K. Ruben- stein, Sylvain Gelly, and Mario Lucie. 2020. On mutual information maximization for representa- tion learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Jean-Philie Philippe Vert, Koji Tsuda, Bernhard Schölkopf, B. Schölkopf K. Tsuda, and Vert J.P. 2004. A primer on kernel methods.

---

A Appendix-A
Model
MRR@ 10
Recali 1000
Params
Dimension
RetroMAE
41.6
98.8
110M
768
A.1 Statistics of the Dataset
Tiay BEa-HEKTD
25.38
92.55
14M
128
BEyDBER-HPP
98.00
14M
The NLI dataset consists of million pairs and 1,196,755 sentences. After data augmentation, the dataset includes 3,590.265 sentences. The MS MARCO dataset includes 502,939 training queries, 8,841,823 documents, and 6,980 test queries.
Tiny BEKT-BKD
30.12
92.04
14M
312
Tiny BEKT-BBK
38.21
98.20
14M
128
Table 7: Results of the DR task with Retromac as the teacher model.
A.2 Hyparameter search range
student model to closely resemble the teacher model in comparison to other methods. This also confirms the generalizability of the IBKD ap- proach.
Table 5 list the range of each hyparameter we used. We select the final hyparameter on the STS-B dev set and MS MARCO dev set for the STS task and DR task respectively.
Table 5: Hyparameters scarch range.
Hyparameter
range
learning rate
1e-5, 3e-5, le-4, le-4 64, 128, 256
batch size
epoch
3, 10
0.01, 0.05, 0.1, 0.2, 0.5 0.01,0 0.5, 1.0, 2.0 0.1, 0.5, 1.0, 2.0
31
82
A.3 Experiment results with additional teacher model
In this section, present experimental results utilizing additional teacher models for the STS and DR tasks. For the STS task, we selected SimCSE- BERTlarge (Gao al., 2021) as our teacher model, while for the DR task, we chose RetroMAE (Xiao et al., 2022). As a baseline, we employed the HPD method. The corresponding results are provided in Table 6 for the STS task and Table 7 for the DR task.
Model
Avg
Params
Dimension
SimCSF-BERTTan BaRTe
82.21
330M
1024
TinyBERT-HPD
80.24
14M
128
BE
80.88
14M
128
BERBBERTIBKKD
81.02
14M
312
TinyBERT-IBKDD 81.43
14M
312
Table 6: Experiment resulst for the STS task with BERT-large as the teacher model.
The additional experiments above indicate that different teacher models IBKD consistently demonstrate performance, effectively enabling the


Reviews:
[{"review_id": "C9POzRZOGp", "reviewer": "Reviewer_QEfC", "summary": "", "questions": "Have you considered applying your method to other representation models beyond Roberta and CoCondenser to verify its generalizability?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 5, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\nWe are very grateful for your evaluation of our paper and the time you've dedicated to providing your feedback. In the subsequent sections, we endeavor to respond to your comments and address the concerns you have raised.\n\n\n> Limited Variety: The proposed method's application is primarily demonstrated on Roberta and CoCondenser. This limited variety in implementation may raise questions about the framework's generalizability. It's unclear whether the method would be as effective if applied to other representation models.\n\nGiven your concern about the generalization of our proposed method, we have chosen additional models as teacher models to validate the generalization ability of our method. It should be noted that, apart from replacing the teacher model, all other hyperparameters remain consistent with the experimental parameters in Table-1 of the original paper. Due to time constraints, we have only selected the HPD method as the baseline for validation.\n\nWe employ the Simcse-BERT-large[1] model as the teacher model for the STS task. We use the term \"Model-IBKD\" to denote a model that has undergone the distillation training stage, and \"Model-IBKD$_{ft}$\" to denote the model after the fine-tuning stage's training. The results are presented below. \n\n| Model                 | STS12 | STS13 | STS14 | STS15 | STS16 | STS-B | SICK-R |  Avg  | Params | Dimension |\n|-----------------------|-------|-------|-------|-------|-------|-------|--------|-------|-------|-------| \n| Simcse-BERT-large     | 75.78 | 86.33 | 80.44 | 86.06 | 80.86 | 84.87 | 81.14  | 82.21 | 330M | 1024 |\n| TinyBERT-HPD          | 72.43 | 83.10 | 77.92 | 84.68 | 80.38 | 83.16 | 79.99  | 80.24 | 14M | 128 |\n| TinyBERT-HPD$_{ft}$       | 73.00 | 83.25 | 79.08 | 84.74 | 81.31 | 83.87 | 80.89  | 80.88 | 14M | 128 |\n| TinyBERT-IBKD         | 74.11 | 83.13 | 78.98 | 85.10 | 81.31 | 84.01 | 80.51  | 81.02 | 14M | 312 |\n| TinyBERT-IBKD$_{ft}$      | 74.48 | 83.39 | 79.40 | 85.28 | 81.80 | 84.44 | 81.21  | 81.43 | 14M | 312 |\n\n\nFor the DR task, we choose another competitve pretrained BERT model RetroMAE[2] as the teacher model. The results are shown below:\n\n| Model             | MRR@10 | Recall@1000 | Params | Dimension |\n|-------------------|--------|-------------|------------|------------| \n| RetroMAE          | 41.6   | 98.8        | 110M | 768 |\n| TinyBERT-HPD      | 25.38  | 92.55       | 14M  | 128  |\n| TinyBERT-HPD$_{ft}$   | 36.93  | 98.04       | 14M  | 128 |\n| TinyBERT-IBKD | 30.12  | 92.04       | 14M | 312 |\n| TinyBERT-IBKD$_{ft}$  | 38.21  | 98.20       | 14M  | 128 |\n\nThe additional experiments above indicate that different teacher models of IBKD consistently demonstrate performance, effectively enabling the student model to closely resemble the teacher model in comparison to other methods. This also confirms the generalizability of the IBKD approach.\n\n> Insufficient experiments: The experimentation conducted in the paper could be more comprehensive. The authors predominantly use a single teacher model for each dataset, which may not fully demonstrate the effectiveness of their method. Including at least one more teacher model in each dataset could provide a more robust validation of the method's efficacy.\n\nWe will add all the above experiment results in the revised version and discussed with them in detail.\n\n> Figure 1 currently suggests that the teacher embeddings encompass all task-relevant information from the original input, which might oversimplify the knowledge distillation process. Authors are suggested to modify Figure 1 to illustrate partial overlap between the teacher embeddings and the task-relevant information, indicating the potential information loss from the teacher model.\n\n\nThank you for your valuable suggestion. The main aim of Figure 1 is to emphasize the problem of over-fitting in the conventional distillation method. To effectively convey this message, we have assumed an ideal scenario where the teacher model incorporates all the required task-related information. However, we recognize that this simplification may have led to some misunderstandings.\nTo address this concern, we will revise Figure 1 to illustrate a partial overlap between the teacher embeddings and the task-relevant information. This adjustment will provide a more precise representation of the real situation and reduce any potential confusion.\n\nThanks again for your comments.\n\n**Reference**:\n\n[1] Simple Contrastive Learning of Sentence Embeddings\n\n[2] RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Good writing and easy to follow\n- It is an inspiring idea to introduce the Information Bottleneck into the knowledge distillation\n- The experiments conducted provide robust evidence for the effectiveness of the proposed method in knowledge distillation", "reasons_to_reject": "- Limited Variety: The proposed method's application is primarily demonstrated on Roberta and CoCondenser. This limited variety in implementation may raise questions about the framework's generalizability. It's unclear whether the method would be as effective if applied to other representation models.\n- Insufficient experiments: The experimentation conducted in the paper could be more comprehensive. The authors predominantly use a single teacher model for each dataset, which may not fully demonstrate the effectiveness of their method. Including at least one more teacher model in each dataset could provide a more robust validation of the method's efficacy.", "questions_to_authors": "Have you considered applying your method to other representation models beyond Roberta and CoCondenser to verify its generalizability?", "paper_topic_and_main_contributions": "The paper proposes the IBKD method for making pre-trained language models (PLMs) more accessible by distilling them into smaller representation models. By implementing the Information Bottleneck (IB) principle, this method reduces information loss during the distillation process and concentrates on the most informative features of the PLMs. A standout aspect of this approach is the addition of a novel regularization term to counteract overfitting, as well as a well-devised implementation strategy for the IB principle. Experimental results underscore the method's effectiveness, demonstrating substantial performance improvements across a variety of text classification tasks.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "Figure 1 currently suggests that the teacher embeddings encompass all task-relevant information from the original input, which might oversimplify the knowledge distillation process. Authors are suggested to modify Figure 1 to illustrate partial overlap between the teacher embeddings and the task-relevant information, indicating the potential information loss from the teacher model."}}, {"review_id": "VGJz7O25hs", "reviewer": "Reviewer_BBAJ", "summary": "", "questions": "A. My main question aligns with the Reasons to Reject, that is, how to use formula (7) to calculate the RBF kernel between two texts?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you for your insightful comments and suggestions. We have carefully reviewed your feedback and have made the following detailed responses:\n\n> The process of estimating mutual information with HSIC is not detailed enough, especially on how to use formula (7) to calculate the RBF kernel between two texts?\n\nWe apologize for the insufficient explanation regarding the calculation process of HSIC. In the revised version, we will address this concern by providing a clearer explanation of the process.\n\nIn our original paper, we introduced formula (6) as the method to calculate the HSIC between the input $X$ and student embedding $S$. To compute HSIC, it is necessary to calculate the kernel Gram matrices $K_x ∈ R^{n \\times n}$ for $X$ and Ks ∈ R^{n x n} for S. Here, $K_{x_{ij}} = k(x_i, x_j)$ and $K_{s_{ij}} = k(s_i, s_j)$, where $x_i$ represents the input representation of the i-th text, and $s_i$ represents the output of the student model for the i-th text. The function $k(\\cdot, \\cdot )$ can be any kernel function. In this study, we utilize the RBF kernel, as described in formula (7), to calculate $k(x_i, x_j)$.\n\n> Though the paper focuses on distillation via the mutual information technique, referencing some other works on BERT distillation would be beneficial:  \\\n> 1. BERT-of-Theseus: Compressing BERT by Progressive Module Replacing. ACL 2020. \\\n> 2. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter.\n\nWe will cite these papers you mentioned and discussed them in the related work section. \n\nThanks again for your comments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "A. The model design is reasonable, including the use of the information bottleneck and the two-stage training method.\nB. The effectiveness of the model is well demonstrated by extensive experiments, and the authors provide detailed parameter settings.\n", "reasons_to_reject": "A. The process of estimating mutual information with HSIC is not detailed enough, especially on how to use formula (7) to calculate the RBF kernel between two texts?\n", "questions_to_authors": "A. My main question aligns with the Reasons to Reject, that is, how to use formula (7) to calculate the RBF kernel between two texts?", "paper_topic_and_main_contributions": "The paper optimizes Text Representation Distillation using the Information Bottleneck Principle. The work falls under computationally-aided linguistic analysis. The authors argue that traditional distillation methods based on mutual information retain some noise information from the teacher model, which is unnecessary for downstream tasks. Therefore, they adopt the Information Bottleneck strategy to constrain the mutual information between the student model representation and the input text to be as small as possible, while distilling information from the teacher model to the student model. Moreover, they propose a two-stage training method to optimize the model. The first stage is distillation learning, and the second stage is model fine-tuning on supervised data. The main contributions of this work are: 1) the application of the information bottleneck theory to language model distillation and the proposal of a two-stage training method to improve model performance; 2) thorough experimental validation of the model's effectiveness, along with the publication of detailed descriptions and key parameters.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "N/A"}, "extra_text": {"Missing_References": "Though the paper focuses on distillation via the mutual information technique, referencing some other works on BERT distillation would be beneficial:\n1) BERT-of-Theseus: Compressing BERT by Progressive Module Replacing. ACL 2020. \n2) DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. \n"}}, {"review_id": "PUf7VMXrNE", "reviewer": "Reviewer_gS4Y", "summary": "", "questions": "A.\tWhy were STS and DR selected as downstream tasks to evaluate distillation performance?\nB.\tWhat data was used for unsupervised learning stage?\nC.\tIt is not clear whether the teacher model was fine-tuned for the STS and DR tasks or not.\n", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, \nWe are very grateful for your recognition of our efforts. We have carefully considered your thoughtful comments and we'd like to offer some clarifications as well as solutions that aim to address the concerns you raised.\n\n> It is stated that “simply maximizing I(T,S) is prone to over-fitting. However, IBKD can effectively filter out task-irrelevant information and retains only the task-relevant information, thus improving the generalization of the student model’s representation”. However, there is no result showing that the proposed method can improve generalization ability.\n\nIt is important to clarify that in the task of text representation model distillation, the generalization ability of the student model is typically determined by the performance of the student model on the test dataset, rather than simply fitting the representation of the teacher model. Alternatively, we can view this as the student model maximizing its fit to the training dataset, as the teacher model is trained on this dataset and can be seen as another form of the training dataset. \n\nTo validate our claims, we selected two widely recognized downstreaming tasks of text representation, namely senmatic text similarity (STS) and dense passage retrieval (DR). Our experimental results demonstrate that traditional model distillation methods, particularly those that directly approximate the teacher and student model representations using the Mean Squared Error (MSE) loss function, significantly degrade performance on the test set and undermine the generalization of the student model. However, the IBKD method proposed in this paper effectively mitigates this loss of generalization performance.\n\nIn the revised version, we will provide a more detailed explanation to elucidate the improvement in model generalization performance.\n\n> The range of tasks used for evaluation is narrow. When assessing distillation performance, many papers use benchmark datasets like GLUE, which evaluate the model's ability in various aspects.\n\n> Why were STS and DR selected as downstream tasks to evaluate distillation performance? \n\nIn this paper, our main objective was to develop a distillation method for the text representation model. Building upon previous work, we selected two downstream tasks that are most representative of text representation to evaluate the performance of our model. As for the GLUE benchmark you mentioned, it is commonly used by researchers to assess the overall performance of pre-trained language models. Typically, this evaluation process involves separate training for each individual task, which is not suitable for the context of our paper focusing on the representation model. However, it is worth noting that our proposed distillation method has the potential to be extended to the distillation process of the entire pre-training base model or other task models, rather than just limited to the distillation of the representation model. We leave further exploration for future studies.\n\n\n> What data was used for unsupervised learning stage\n\nThe training data used for the unsupervised learning stage was discussed in Section 4.3 of the original paper. In the case of the STS task, we employed an augmented Natural Language Inference (NLI) dataset. This dataset comprises the original NLI dataset, supplemented with extra data created by applying WordNet substitution and back translation to every instance. As for the DR task, we utilized all passages and training queries from the MS MARCO Passage Dataset during the initial training stage.\n\n> It is not clear whether the teacher model was fine-tuned for the STS and DR tasks or not.\n\nThe teacher model for each task undergoes fine-tuning using task-specific supervised data. The coCondenser model is trained on the MSMARCO dataset, while Simcse-Roberta is fine-tuned using the NLI dataset. \n\nWe will clarify all these details above in the revised version and fix all typos.  Thanks again for your comments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "This method outperformed other distillation methods on two downstream tasks (STS, DR).\nThis paper showed that the contrastive learning loss and HSIC method were effective for performance improvement.\n", "reasons_to_reject": "It is stated that “simply maximizing I(T,S) is prone to over-fitting. However, IBKD can effectively filter out task-irrelevant information and retains only the task-relevant information, thus improving the generalization of the student model’s representation”. However, there is no result showing that the proposed method can improve generalization ability. \n\nThe range of tasks used for evaluation is narrow. When assessing distillation performance, many papers use benchmark datasets like GLUE, which evaluate the model's ability in various aspects.\n", "questions_to_authors": "A.\tWhy were STS and DR selected as downstream tasks to evaluate distillation performance?\nB.\tWhat data was used for unsupervised learning stage?\nC.\tIt is not clear whether the teacher model was fine-tuned for the STS and DR tasks or not.\n", "paper_topic_and_main_contributions": "This paper proposes a knowledge distillation method in text representation called IBKD that aims to address the performance gap between large-scale pre-trained models and smaller ones. Inspired from the Information Bottleneck (IB) principle, the method maximizes mutual information between the final representation of the teacher and student model and minimizes mutual information between the representation of the student model and input. The paper also introduces the contrastive learning loss and HSIC method to reduce the computational cost of mutual information. ", "excitement": 3, "reproducibility": 5, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "Is it correct that the dimension of TinyBERT-IBKD-128 is 312? It seems like a typo."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes the IBKD method in order to make pre-trained language models (PLMs) more accessible by distilling them into smaller representation models. The topic studied is an interesting topic and the proposed method based on information bottleneck is meaningful. Reviewers have concerns on the insufficient experiments or missing details. During the rebuttal period, the authors have replied to these concerns with more details and experimental results, which lead to overall positive scores from reviewers.
Decision: Accept-Main