SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} = \mathcal {L} _ {g} + \alpha \mathcal {L} _ {t} + \beta \mathcal {L} _ {s} \tag {10}
$$
Here,  $\mathcal{L}_g$  denotes the task training loss, while  $\alpha$ ,  $\beta$  are weighting factors used to balance the task training loss and the relation losses.
As discussed in Section 3.3, different distance functions like Kullback-Leibler (KL) divergence and Mean Squared Error (MSE) are commonly used to measure the similarity between the teacher and the student models. However, these functions have an 'exact match' property, which means that the distance is zero if and only if the student and teacher features are exactly the same. This
Model	#Params	Speedup	SST-2	MNLI-m	QNLI	QQP	RTE	SST-B	MRPC	CoLA	AVG
BERTBASE	110M	×1.0	93.4	84.5	91.5	72.3	66.8	85.2	88.3	52.8	79.3
BERTSMALL6	66M	×2.0	90.7	81.2	87.9	69.4	64.3	79.8	83.7	41.4	74.8
BERT-PKD6	66M	×2.0	92.0	81.5	89.0	70.7	65.5	81.6	85.0	43.5	76.1
DistilBERT6	66M	×2.0	92.5	82.6	88.9	70.1	58.4	81.3	86.9	49.0	76.2
TinyBERT6	66M	×2.0	92.1	82.8	89.8	71.2	70.0	83.9	88.0	51.1	78.6
MinILM6	66M	×2.0	92.0	83.0	91.1	71.4	70.8	84.2	88.5	49.2	78.8
MinILMv26	66M	×2.0	92.4	83.4	90.0	71.5	71.3	84.5	88.6	51.8	79.2
Ours	66M	×2.0	92.8	83.8	91.3	72.0	71.7	84.8	89.1	52.0	79.6
RoBERTaBASE	125M	×1.0	95.3	87.2	93.2	73.8	72.7	88.4	90.1	62.0	82.8
RoBERTaSMALL6	82M	×2.0	92.3	83.1	90.4	72.1	68.4	86.8	87.5	54.1	79.3
MinILMv26	82M	×2.0	93.5	84.3	91.6	72.8	72.1	87.5	88.2	57.8	81.0
Ours	82M	×2.0	93.8	85.6	92.0	73.5	72.5	88.3	89.6	60.3	81.9
DistilGPT2	82M	×2.3	90.7	81.6	87.9	66.8	68.3	79.6	87.9	39.4	75.3
Ours	82M	×2.3	92.0	83.4	88.5	70.6	70.2	81.6	88.4	42.3	77.1
Table 1: Results of the proposed method on the test sets of GLUE. We use the metric of Matthews correlation for CoLA, Pearson-Spearman correlation for STS-B, and accuracy for other datasets. Following previous works Sun et al. [2019a], we also report the average score of these eight tasks (the "AVG" column). The speedup is in terms of the BERT<sub>BASE</sub> and RoBERT<sub>BASE</sub> inference time and evaluated on a single GPU with a single input of 64 or 128 length. The fine-tuning results are an average of 4 runs.
can be written as  $\Gamma(X, Y) = 0$  when  $X = Y$ . Unlike KL divergence or MSE, PLC is invariant under positive linear transformations. This means that even if a positive linear transformation is applied to one or both of the features, the correlation coefficient remains the same. In other words,  $\Gamma(X, Y) = 0$  if  $Y = \alpha X + \beta$ , where  $\alpha > 0$  and  $\beta$  are constants. This property makes PLC a more flexible choice for knowledge distillation, as it allows the student model to learn from the teacher model in a less restrictive way. The detailed mathematical justification for this property is provided in the supplementary materials.
We next delve into the relationship between the Pearson linear correlation (PLC), Kullback-Leibler (KL) divergence and Mean Squared Error (MSE). For KL divergence, the normalized features are transformed into a probability distribution, and then we minimize the discrepancy between the softened probabilities of the teacher and student models.
$$
\mathcal {L} _ {K L} (\hat {\mathcal {F}} _ {T}, \hat {\mathcal {F}} _ {S}) = \tau^ {2} \sum_ {m} \phi_ {t} (\hat {\mathcal {F}} _ {T}; \tau) \log \frac {\phi_ {t} (\hat {\mathcal {F}} _ {T} ; \tau)}{\phi_ {s} (\hat {\mathcal {F}} _ {S} ; \tau)} \tag {11}
$$
Here,  $\tau$  denotes a temperature parameter used to adjust the softness degree in the distributions. Assuming the value of  $\tau$  is significantly large compared to the magnitude of the normalized features, and  $\hat{\mathcal{F}}_S, \hat{\mathcal{F}}_T$  are drawn from a standard normal distribution, we can derive the gradient of  $\mathcal{L}_{KL}$  with respect to the normalized feature  $\hat{\mathcal{F}}_S^i$  as follows:
$$
\frac {\partial \mathcal {L} _ {K L}}{\partial \hat {\mathcal {F}} _ {S}} \approx \frac {1}{M} \left(\hat {\mathcal {F}} _ {S} - \hat {\mathcal {F}} _ {T}\right) = \frac {\partial \mathcal {L} _ {M S E}}{\partial \hat {\mathcal {F}} _ {S}} \tag {12}
$$
Further, given that  $\frac{1}{m - 1}\sum_{i}\hat{\mathcal{F}}_S^2 = 1$  and  $\frac{1}{m - 1}\sum_{i}\hat{\mathcal{F}}_T^2 = 1$ , we can reformulate the MSE loss as follows:
$$
\mathcal {L} _ {M S E} \left(\hat {\mathcal {F}} _ {S}, \hat {\mathcal {F}} _ {T}\right) = \frac {1}{2 M} \sum \left(\hat {\mathcal {F}} _ {S} - \hat {\mathcal {F}} _ {T}\right) ^ {2} \approx 1 - \rho \left(\mathcal {F} _ {S}, \mathcal {F} _ {T}\right) \tag {13}
$$
This equivalence clarifies the intrinsic connection between KL divergence, MSE, and PLC of normalized feature distributions. The detailed derivations and empirical results to support this claim are provided in the supplementary material.
4.1 Distillation Setup
We evaluate the efficacy of the proposed FCD on 8 out of 9 tasks from the General Language Understanding Evaluation (GLUE) Wang et al. [2019] benchmark, which consists of 2 single-sample
Figure 3: Comparison of distillation results of student models using BASE-size and LARGE-size teachers (BERT and RoBERTa) shown in the left part. The right part shows ablation of token-level and sample-level relations.
(CoLA and SST-2) and 5 sample-pair (MRPC, QQP, MNLI, QNLI and RTE) classification tasks, and 1 regression task (STS-B). Following previous worksSun et al. [2019a], we use the same metrics as the GLUE benchmark for evaluation. In order to verify the effectiveness and robustness of our method, we distill teacher models with different architectures, model sizes. Concretely, we consider encode-based models including BERT<sub>BASE</sub>, BERT<sub>LARGE</sub>, RoBERT<sub>BASE</sub> and RoBERT<sub>LARGE</sub> to ensure a fair comparison with of a wide range of prior works. Moreover, we also explore the compression of decoder-based models by employing our distillation method to improve fine-tuning of DistilGPT2 Wolf et al. [2020], which is rarely investigated in most previous works.
The process of distilling pretrained Language Models (LM) generally comprises two stages: task-agnostic distillation and task-specific distillation. Task-agnostic distillation involves a pre-training process on a large-scale dataset. However, this stage can be costly and time-consuming, particularly for scenarios with limited computational resources. In contrast, task-specific LM distillation proves to be effective and considerably more economical compared to pre-training. As an increasing number of pretrained LM models are becoming publicly available through resources such as the HuggingFace Transformers library Wolf et al. [2020], directly leveraging these models can save substantial time and computational resources compared to training from scratch. Given these considerations, our focus in this work is on task-specific knowledge distillation. However, it is important to note that our method is not limited to task-specific distillation and can be readily applied in task-agnostic scenarios as well. Specifically, we first fine-tune the pretrained teacher models on a specific task. Following this, the corresponding student model is initialized with the teacher model using the LayerDropping method Sajjad et al. [2020]. Subsequently, we perform distillation with FCD. We employ a grid search algorithm on the development set to tune the hyper-parameters. Specifically, we trained the student model for 3, 5 and 10 epochs, using a batch size of 32. The learning rates we experimented included  $2e - 5$ ,  $1e - 5$  and  $5e - 6$ . For the CoLA task, we extended the training steps to 25 epochs. The parameters  $\alpha$  and  $\beta$  from the distillation loss are tuned from  $\{0.1, 0.2, 0.4, 1\}$ , a choice guided by maintaining the different components of the loss in the same order of magnitude. We adopt a cosine decay schedule with a warm-up phase of 5 epochs and utilize the AdamW optimizer with a weight decay of 0.5. The maximum sequence length is set to 64 for single-sample tasks, and 128 for sequence pair tasks.
We start by comparing our proposed method with several KD baselines, including DistilBERT Sanh et al. [2019], TinyBERT Jiao et al. [2019], BERT-PKD Sun et al. [2019a], MiniLM Wang et al. [2020b], MiniLMv2 Wang et al. [2020a]. Similar to previous studies, we distill a 12-layer base model into a 6-layer student model with only about  $60\%$  parameters and  $2x$  inference speedup. In order to evaluate the impact of knowledge distillation, we also report the results of  $\mathrm{BERT}_{\mathrm{SMALL6}}$  and  $\mathrm{RoBERTa}_{\mathrm{SMALL6}}$ . These smaller models are obtained using the LayerDropping method Sajjad et al. [2020], wherein the strategy of dropping the top layer has been demonstrated to be a strong baseline. Consequently, we drop the top 6 layers of the base model and fine-tune it without using knowledge distillation. For fair comparisons, we fine-tune the released models and evaluate the result on the test set of GLUE without resorting to data augmentation strategy Jiao et al. [2019]. The results
Table 2: Comparison of using different distillation loss functions.
Method	MNLI-m	QQP	RTE	Average
FCD (KL)	83.3	71.2	70.7	75.1
FCD (MSE)	83.6	71.6	71.2	75.5
FCD (Pearson)	83.8	72.0	71.7	75.8
Table 3: Distillation results with different layer selecting schemes.
Scheme	layers	MNLI-m	CoLA
Top	{7,8,9,10,11,12}	82.9	50.7
Uniform	{2,4,6,8,10,12}	83.8	52.0
Bottom	{1,2,3,4,5,6}	83.3	51.5
from these 6-layer student models are summarized in Table 1. The top group of models denotes the uncased version of base-size BERT used as the teacher model. Notably, our model outperforms all compared models by a large margin. Similar trends are observed in the middle group, where base-size RoBERTa is utilized as the teacher model. Our model surpasses the MiniLMv2 by  $0.4\%$  accuracy on RTE,  $0.5\%$  F1 on MRPC, and  $0.3\%$  Spearman correlation on STS-B. Furthermore, we conduct experiments on a decode-base model DistilGPT2. In this setting, we employ the GPT-2 model with 12 layers and 768 hidden size as the teacher model. Noatably, our proposed method outperforms the original KD method by an average of  $1\%$ , underlining the efficacy of our approach in decoder-based model. As illustrated in the left part of Figure 3, student models that are distilled from large-size teacher models achieve further improvements. Moreover, this performance gain increases with the capacity of teacher models, thereby demonstrating the effectiveness of our proposed method across different sizes of pretrained Transformer models.
Effect of different components In this study, we introduce two distinct types of relations: token-level and sample-level relations. To verify the effectiveness of each, we conduct experiments using  $\mathcal{L}_t$  and  $\mathcal{L}_s$  to investigate their respective influences on the student model. As depicted in the right part of Figure 3, each component within the distillation loss independently contributes to the enhancement of the final performance. Moreover, a further boost in performance is observed when these components are combined. We noticed that the performance degradation on SST-2 is more substantial compared to other tasks without token-level relation. We speculate that the token-level relation is particularly important for this single sentence binary classification task.
Effect of different distillation loss functions Here, we compare our proposed Pearson Linear Correlation (PLC) with the Mean Squared Error (MSE) and Kullback-Leibler (KL) divergence, which are widely-used loss functions. To ensure a fair comparison, we tune the distillation loss weight for both MSE and KL. The comparative results across three tasks are presented in Table 2, which demonstrates that adopting Feature Correlation Distillation (FCD) with PLC consistently yields higher performance compared to the FCD combined with MSE and KL. This indicates that the more flexible Pearson correlation might serve as a more suitable metric for measuring relations within the FCD.
Effect of different layer selection strategies Apart from the type of knowledge used for distillation, the selection of layers significantly affects the overall performance of distillation. We study the impact of three distinct layer selection strategies: uniform, top, and bottom, and compare their respective performances. Specifically, we utilize  $\mathrm{BERT}_{\mathrm{BASE}}$  as the teacher model and a  $6 \times 768$  model as the student model. The number of selected layers is set to 6. The results are reported in Table 3. For  $\mathrm{BERT}_{\mathrm{BASE}}$ , using the uniform layer selection strategy yields superior performance compared to the other strategies. This finding highlights the crucial role of both the head and tail layers of student models in the distillation process.
In this paper, we introduce Feature Correlation Distillation (FCD), a novel and effective method for distilling large Transformer-based Pretrained Language Models (PLMs). Our approach simultaneously models both token-level and sample-level relations derived from the features of the Transformer block. Moreover, we propose a correlation-based distillation loss to enhance the performance of the model distillation process. We also provide a theoretical interpretation of our proposed Pearson
linear correlation formulation, offering a deeper understanding of its underlying operation and implications. Through extensive experiments on the GLUE tasks, our distilled smaller language models consistently outperformed existing knowledge distillation methods across a variety of architectures while significantly reducing both the model size and inference time. With its simplicity and strong performance, we hope our approach can serve as a solid baseline for future research.
This research was supported by Ant Group. We are grateful to the five anonymous reviewers for their insightful suggestions and comments, which significantly enhanced the quality of this paper.
6.1 Implementation of FCD
The implementation details of FCD are presented in Figure 4. Utilizing the output features from both the student and teacher models, denoted as  $F_{s}$  and  $F_{t}$  respectively, our method ensures ease of implementation.
import torch nn as nn   
from torch nn import functional as F   
def cosineSimilarity(x,y,eps=1e-8): return (x  $\ast y$  ).sum(1) / (x(norm(dim=1)  $\ast$  y(norm(dim=1)  $^+$  eps)   
def pearson_corrlation(x,y,eps=1e-8): return cosine_similarity(x-x.mean(-1).unsqueeze(-1), y-y.mean(-1).unsqueeze(-1),eps)   
def plc_loss(x,y): return 1 - pearson_corrlation(x,y).mean()   
def token Relation_loss(y_s,y_t): R_s=y_s.bmm(y_s.transpose(-1,-2)) R_t=y_t.bmm(y_t.transpose(-1,-2)) R_s=R_s.view(R_s.size(0), -1) R_t=R_t.view(R_t.size(0), -1) loss  $=$  plc_loss(R_s,R_t) return loss   
def sample Relation_loss(y_s,y_t): y_s=y_s.permute(1,0,2) y_t=y_t.permute(1,0,2) loss  $=$  token relation_loss(y_s,y_t) return loss   
class FCDLoss(nnModule): def init_(self, alpha, beta): super(FCDLoss,self).init_(self.alpha  $=$  alpha self.beta  $=$  beta   
def forward(self,F_s,F_t): F_s  $=$  F normalize(F_s,dim=-1) F_t  $=$  F normalize(F_t,dim=-1) loss_token  $=$  token_relation_loss(F_s,F_t) loss_sample  $=$  sample_relation_loss(F_s,F_t) kd_loss  $=$  self.alpha \* loss_token  $^+$  self.beta \* loss_sample return kd_loss
Figure 4: The PyTorch implementation of FCD.
6.2 Proofs
Invariance of Pearson's Correlation under Positive Linear Transformation. Let's consider two random variables  $X$  and  $Y$ . A positive linear transformation on  $X$  and  $Y$  can be formulated as  $X' = aX + b$  and  $Y' = cY + d$ , where  $a \times c > 0$  and  $b, d$  are arbitrary constants. Applying these transformations to the means of  $X$  and  $Y$ , we derive  $\mu_{X'} = a\mu_X + b$  and  $\mu_{Y'} = c\mu_Y + d$ . By substituting the transformed variables and their corresponding means into Equation 6:
$$
\begin{array}{l} \rho^ {\prime} (X, Y) = \frac {\sum ((a X _ {i} + b) - (a \mu_ {X} + b)) ((c Y _ {i} + d) - (c \mu_ {Y} + d))}{\sqrt {\sum ((a X _ {i} + b) - (a \mu_ {X} + b)) ^ {2}} \sqrt {\sum ((c Y _ {i} + d) - (c \mu_ {Y} + d)) ^ {2}}} \\ = \frac {\sum (a (X _ {i} - \mu_ {X})) (c (Y _ {i} - \mu_ {Y}))}{(a \sqrt {\sum (X _ {i} - \mu_ {X}) ^ {2}}) (a \sqrt {\sum (Y _ {i} - \mu_ {Y}) ^ {2}})} = \rho (X, Y) \\ \end{array}
$$
Relationship among PLC, KL divergence and MSE In Equation 11,  $\phi (\cdot)$  represents the softmax function, while  $\tau$  denotes a temperature parameter controlling the softness of the distributions.
$$
\frac {\partial \mathcal {L} _ {K L}}{\partial \hat {\mathcal {F}} _ {S} ^ {i}} = \tau (\phi_ {s} (\tau) - \phi_ {t} (\tau)) = \tau \left(\frac {\exp \left(\hat {\mathcal {F}} _ {S} ^ {i} / \tau\right)}{\sum_ {j = 1} ^ {m} \exp \left(\hat {\mathcal {F}} _ {S} ^ {j} / \tau\right)} - \frac {\exp \left(\hat {\mathcal {F}} _ {T} ^ {i} / \tau\right)}{\sum_ {j = 1} ^ {m} \exp \left(\hat {\mathcal {F}} _ {T} ^ {j} / \tau\right)}\right), \tag {15}
$$
Assuming  $\tau$  is significantly large compared to the magnitude of the normalized features and both  $\hat{\mathcal{F}}_S^i$  and  $\hat{\mathcal{F}}_T^i$  are drawn from a standard normal distribution. In this case, the term  $\hat{\mathcal{F}}^i / \tau$  becomes quite small, allowing us to approximate  $\exp (\hat{\mathcal{F}}^i /\tau)$  as  $1 + \hat{\mathcal{F}}^i /\tau$ . This simplification leads to an
(c)
Figure 5: Visualization of the token-level relational features under different normalization functions. (a) Pre-Norm. (b)  $\ell_2$ -Norm. (c) Layer-Norm. (d) Softmax-Norm.
approximation of the gradient in Equation 15:
$$
\frac {\partial \mathcal {L} _ {K L}}{\partial \hat {\mathcal {F}} _ {S} ^ {i}} \approx \tau \left(\frac {1 + \hat {\mathcal {F}} _ {S} ^ {i} / \tau}{M + \sum_ {j = 1} \hat {\mathcal {F}} _ {S} ^ {j} / \tau} - \frac {1 + \hat {\mathcal {F}} _ {T} ^ {i} / \tau}{M + \sum_ {j} \hat {\mathcal {F}} _ {T} ^ {j} / \tau}\right) \tag {16}
$$
Given that the sums  $\sum_{j}\hat{\mathcal{F}}_{S}^{j}$  and  $\sum_{j}\hat{\mathcal{F}}_{T}^{j}$  are both zero, Equation 16 simplifies further to:
$$
\frac {\partial \mathcal {L} _ {K L}}{\partial \hat {\mathcal {F}} _ {S} ^ {i}} = \frac {1}{M} (\hat {\mathcal {F}} _ {S} - \hat {\mathcal {F}} _ {T}) = \frac {\partial \mathcal {L} _ {M S E}}{\partial \hat {\mathcal {F}} _ {S} ^ {i}} \tag {17}
$$
Moreover, considering  $\frac{1}{m - 1}\sum_{i}\hat{\mathcal{F}}_S^2 = 1$  and  $\frac{1}{m - 1}\sum_{i}\hat{\mathcal{F}}_T^2 = 1$ , we can reformulate the MSE as follows:
$$
\begin{array}{l} \mathcal {L} _ {M S E} (\hat {\mathcal {F}} _ {S}, \hat {\mathcal {F}} _ {T}) = \frac {1}{2 M} \sum (\hat {\mathcal {F}} _ {S} - \hat {\mathcal {F}} _ {T}) ^ {2} \\ = \frac {1}{2 M} \left((2 M - 2) - 2 \sum_ {i = 1} ^ {m} \hat {\mathcal {F}} _ {S} \hat {\mathcal {F}} _ {T}\right) \tag {18} \\ = \frac {2 M - 2}{2 M} (1 - \rho (\mathcal {F} _ {S}, \mathcal {F} _ {T})) \approx \mathcal {L} _ {P L C} (\mathcal {F} _ {S}, \mathcal {F} _ {T}) \\ \end{array}
$$
Thus, we demonstrate that minimizing KL divergence between normalized features under a high-temperature limit is equivalent to minimizing the MSE between normalized ones, which is in turn equivalent to maximizing the PLC between the original features.
6.3 More Experiments Results
Effect of different Normalization and Loss Functions Section 3.4 in the main text clarifies the intrinsic relationship between KL divergence, MSE, and PLC. However, the assumption that normalized features follow a Gaussian distribution may not invariably be valid. To investigate the performance of varying normalization and loss functions, we conducted a series of experiments, setting the temperature  $\tau$  to 10 when utilizing KL divergence as the loss function. Table 4 demonstrates that the  $\ell 2$  norm consistently outperforms other normalization functions. Minimizing KL divergence between layer-normalized features in the high-temperature limit can yield results comparable to MSE and PLC. To further underscore the benefits of  $\ell 2$  normalization, we provide a visualization of the pre-normalized and post-normalized token-level relational features in Figure 5. In contrast to  $\ell 2$  normalization, other functions often produce a wider range with larger values, suggesting that directly imitating these normalized features could introduce significant noise, potentially leading to subpar results.
Results on SQuAD v1.1 and v2.0. To further demonstrate FCD's effectiveness, we applied it to the question-answering tasks of SQuAD v1.1 Rajpurkar et al. [2016] and SQuAD v2.0 Rajpurkar et al. [2018]. We framed these tasks as sequence labeling problems, predicting the likelihood of each token being the start or end of an answer span. We employed the F1 metric for both versions of SQuAD.  $\mathrm{BERT}_{\mathrm{BASE}}$  was used as the teacher model, and a  $6\times 768$  model served as the student model. The results, presented in Table 5, indicate that FCD can enhance the student model's performance on both tasks.
Table 4: Results of the loss function combined with different normalization mechanism.
Scheme	MNLI-m	QQP	CoLA	Average
LayerNorm + KL	83.4	71.5	51.4	68.8
LayerNorm + MSE	83.6	71.8	51.5	69.0
PreNorm + PLC	83.3	71.2	51.7	68.7
Softmax + PLC	83.4	71.7	51.5	68.9
\( \ell_2 + \)PLC	83.8	72.0	52.0	69.3
Table 5: Results of baselines and FCD on question answering tasks.
Method	SQuAD 1.1	SQuAD 2.0
BERTBASE	88.7	78.8
DistilBERT6	86.2	69.5
TinyBERT6	87.5	77.7
Ours	88.2	78.4
Limitations. While FCD demonstrates consistent performance improvements across diverse Transformer-based models, its effectiveness may be less pronounced on other architectures such as Recurrent Neural Networks (RNNs). The feature relationships in RNNs are not as explicit as in Transformers, potentially limiting the applicability and impact of FCD.
Societal impacts. The extensive computational resources required to evaluate our proposed method could significantly contribute to carbon emissions, thereby raising sustainability concerns. However, the objective of our approach is to enhance the efficiency of lightweight models through knowledge distillation. This enhancement could ultimately replace heavier models in production settings, resulting in substantial energy savings. Thus, the thorough validation of FCD's efficacy is a necessary trade-off to ensure its potential benefits.