SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kun Huang $^{1}$  Xin Guo $^{1}$  Meng Wang $^{1*}$
{hunterkun.hk, darren.wm, bangzhu.gx}@antgroup.com
Knowledge Distillation (KD) has emerged as a promising approach for compressing large Pre-trained Language Models (PLMs). The performance of KD relies on how to effectively formulate and transfer the knowledge from the teacher model to the student model. Prior arts mainly focus on directly aligning output features from the transformer block, which may impose overly strict constraints on the student model's learning process and complicate the training process by introducing extra parameters and computational cost. Moreover, our analysis indicates that the different relations within self-attention, as adopted in other works, involves more computation complexities and can easily be constrained by the number of heads, potentially leading to suboptimal solutions. To address these issues, we propose a novel approach that builds relationships directly from output features. Specifically, we introduce token-level and sequence-level relations concurrently to fully exploit the knowledge from the teacher model. Furthermore, we propose a correlation-based distillation loss to alleviate the exact match properties inherent in traditional KL divergence or MSE loss functions. Our method, dubbed FCD, presents a simple yet effective method to compress various architectures (BERT, RoBERTa, and GPT) and model sizes (base-size and large-size). Extensive experimental results demonstrate that our distilled, smaller language models significantly surpass existing KD methods across various NLP tasks.
Past few years have witnessed a rapid development of pre-trained language models (PLMs) thanks to their effectiveness across a wide range of natural language processing tasks. Pre-trained language models, such as BERT Devlin et al. [2018], RoBERTa Liu et al. [2019], and GPT-2 Radford et al. [2019], learn contextualized text representations by predicting words given their context using large scale text corpora, and can be fine-tuned with additional task-specific layers to adapt to downstream tasks. However, the excellent capability for various NLP tasks comes at demaning huge resources and large memory footprints. For example, the  $\mathrm{BERT}_{\mathrm{BASE}}$  model contains about 110M parameters and 12 Transformer Vaswani et al. [2017] layers, which prevents these transformer-based models from being finetuned and deployed on resource-constrained devices and real-time applications. Recent studies Kovaleva et al. [2019], Voita et al. [2019] indicate that redundancy exists in the original PLMs. Therefore, a series of attempts Chung et al. [2020], Wu et al. [2020], Wang et al. [2020c], Gordon et al. [2020a], Tang et al. [2019], Aguilar et al. [2019] have been made to review the techniques for effective compression of the pre-trained heavy transformers without compromising the performance, of which knowledge distillation is considered to be a practical paradigm. Typically, knowledge distillation techniques aims at effectively transferring the dark knowledge embedded in a large teacher network to boost the performance of the smaller student network during training. Once trained, this compact
student network can be directly deployed in real-life applications without introducing extra inference time or structure modifications. The essence of knowledge distillation relies on how to formulate and transfer the knowledge from teacher to student. The classic logit-based distillation Hinton et al. [2015] directly mimic the final prediction outputs between the teacher and student via Kullback-Leibler (KL) divergence, which only brings limited performance gain to the student. Besides this vanilla knowledge distillation, many other works Du et al. [2020], Heo et al. [2019], Tian et al. [2019] also try to make use of intermediate representations of the pre-trained teacher transformer. The intermediate layers contain more embedding, supplement richer features, thus allow the student transformer to acquire more information in addition to outputs. Jiao et al. [2019] proposed TinyBERT, which distill the information between multiple intermediate features, including the embedding, self-attention matrices and hidden states of the teacher and student networks via the mean squared error (MSE). However, this usually need the adaption layers to align the mismatching embedding dimensions. Such gap makes it hard for the student to mimic the teacher's feature directly and induces additional training cost as a consequence. MINILM Wang et al. [2020b] employs self-attention heatmaps and and value relations via the KL-divergence loss to deeply mimic teacher's self-attention modules. It leads to a restriction that the number of attention heads of student model has to be the same as its teacher. To solve this problem, MINILMv2 Wang et al. [2020a] first concatenate and then split self-attention vectors of different attention heads according to the desired number of relation heads, which involves more computation of queries, keys, and values in self-attention and consequently leads to suboptimal performance.
Motivated by these observations, we aim to directly model feature relationships between the teacher and student models. In a manner similar to the self-attention mechanism, we first model token relations using the output features of the teacher and student models. This token-level relationship has the capacity to capture long-term dependencies between input tokens and highlight critical tokens essential for linguistic comprehension. While the token-level relation is intuitive, it reflects only one aspect of the feature relationships. We further exploit another important aspect of feature relations, the sample-level, to capture the semantic relationship across a batch of samples, an aspect that has been overlooked in previous works. In this way, both types of relations are combined to complement each other to bring more fine-grained feature knowledge. Therefore, the student model is expected to have superior performance compared to that trained stand-alone. Thanks to the same shape of feature relations between student and teacher, our method offers increased flexibility with respect to the embedding dimension and the number of attention heads. Moreover, we propose a correlation-based loss function to replace the KL divergence and MSE used in traditional KD methods. Specifically, we employ the pearson linear correlation as a novel loss function, relaxing the exact match property typically associated with KL divergence and MSE. The overview of the proposed method is illustrated by Figure 1. To sum up, our main contributions are outlined as follows:
- We directly model feature relationships between teacher and student models, which jointly leveraging token-level and sample-level relations to distill knowledge for the first time.
- We propose a correlation-based loss function using Pearson linear correlation, and theoretically explain that it offers a more flexible alternative to traditional KL divergence and MSE.
- Extensive experiments are conducted with popular variants of PLMs, including BERT, RoBERTa, and GPT on GLUE datasets, and our proposal consistently performs better than existing methods.
2.1 Pretrained Language Models
Pretrained language models are pretrained on large amounts of text corpus, and then fine-tuned on task-specific dataset. BERT Devlin et al. [2018] proposes to use a masked language modeling (MLM) objective to pretrain a deep bidirectional Transformer encoder. RoBERTa Liu et al. [2019] achieves strong performance by training longer steps using large batch size and more text data. Besides those encoder-based models, GPT-2 Radford et al. [2019] is a decoder-based model designed for unidirectional, left-to-right text processing. It predicts the next word in a sequence given the preceding words, which allows it to generate coherent and contextually relevant text. In addition to knowledge distillation, the compression of pretrained language models have been widely explored, ranging from
Figure 1: Overview of the proposed method Feature Correlation Distillation (FCD). To demonstrate the effectiveness of FCD. We introduce the token-level relationship and sample-level relationship to distill the knowledge from teacher to student. The correlation loss based on pearson linear correlation is used to capture the relationship between the teacher and the student (best viewed in color).
unstructured pruning Gordon et al. [2020b], Guo et al. [2019], attention head pruning Michel et al. [2019]; to layer factorization Lan et al. [2019], quantization Zhang et al. [2020], Bai et al. [2020] and dynamic width/dpeth inference Hou et al. [2020]. However, some of the techniques like weight pruning (irregular sparsity) and quantization typically require complex pipelines and can not lead to inference speedup and run-time memory saving directly without dedicated hardware/libraries (e.g. for sparse or low-bit computing operation). By contrast, knowledge distillation has been found to be a simple and much effective model compression technique that allows a relatively simple model to perform tasks almost as accurately as a complex model. Moreover, it can be combined with other compression techniques (i.e., where the student model is a smaller, quantized, or pruned version of the teacher model) to further compress the pre-trained language models.
Knowledge Distillation (KD) is a process that transferring knowledge from a large teacher model to a small student model. It was first proposed by Hinton et al. [2015] and then how to effectively transfer more knowledge has been explored by many subsequent works Romero et al. [2014], Ahn et al. [2019], Park et al. [2019], Tian et al. [2019], Tung and Mori [2019]. The intermediate layers contain much richer representation, thus allow the student transformer to acquire more information in addition to outputs. Tang et al. [2019] distill fine-tuned BERT into an extremely small bidirectional LSTM. Turc et al. [2019] initialize the student with a small pre-trained LM during task-specific distillation. Sun et al. [2019a] introduce the hidden states from every k layers of the teacher to perform knowledge distillation layer-to-layer. Aguilar et al. [2019] further introduce the knowledge of self-attention distributions and propose progressive and stacked distillation methods. Task-specific distillation requires to first fine-tune the large pre-trained LMs on downstream tasks and then perform knowledge transfer. The procedure of fine-tuning large pre-trained LMs is costly and time-consuming, especially for large datasets. MiniBERT Tsai et al. [2019] uses the soft target distributions for masked language modeling predictions to guide the training of the multilingual student model and shows its effectiveness on sequence labeling tasks. DistillBERT Sanh et al. [2019] uses the soft label and embedding outputs of the teacher to train the student. TinyBERT Jiao et al. [2019] and MOBILE-BERT Sun et al. [2019b] further introduce self-attention distributions and hidden states to train the student. For example, MOBILE-BERT employs inverted bottleneck and bottleneck modules for teacher and student to make their hidden dimensions the same. TinyBERT uses a uniform-strategy to map the layers of teacher and student when they have different number of layers, and a linear matrix is introduced to transform the student hidden states to have the same dimensions as the teacher. However, the presence of those extra modules not only adds burden on network complexity but also complicates the training procedure.
In this section, we first describe the transformer architectures and define the distillation target in Section 3.1. Then we introduce the two proposed types of feature relationships: token-level relation
Token-level Relation
Figure 2: Our method proposes to maintain the token-level and sample-level relations between student and teacher. Token-level relation: relation between the tokens within each sample of teacher and student. Sample-level relation: relation between the samples on a specific token.
Sample-level Relation
in Section 3.2 and sample-level relation in Section 3.2. At last, we elaborate on the correlation-based loss in Section 3.3 and provide a theoretical analysis in Section 3.4.
The majority of pre-trained language models are based on Transformer architectures, which are composed of a stack of Transformer blocks. We first tokenize the input sample into a sequence of tokens and pack them together with special tokens such as [SEP] and [CLS]. These tokens are then projected into token embeddings and fed into transformer blocks. Each Transformer block consists of a Multi-Head Attention (MHA) and a Feed-Forward network (FFN). Layer normalization (LN) Ba et al. [2016] and residual connection He et al. [2016] are integrated around each of these two sub-blocks. Suppose a teacher model  $T$  and a student model  $S$ , the model takes the feature  $\mathcal{F}_{l - 1}$  of the  $l$ -th Transformer block as input, In multi-head attention, heads are computed in parallel to get the final output, which can be formulated as:
$$
\mathcal {A} _ {l} = \operatorname {A t t n} \left(\mathcal {F} _ {l - 1}, Q _ {l}, K _ {l}\right)
$$
$$
\mathcal {H} _ {l} = \sum_ {H} \mathcal {A} _ {l} V _ {l} \tag {1}
$$
$\mathcal{A}_l \in \mathbb{R}^{H \times N \times N}$  is the attention matrix, where  $H$  denotes the number of heads and  $N$  the sequence length of the input. It is calculated as scaled dot-product between  $Q_l$  and  $K_l$  and then apply softmax operation on the each column of matrix  $\mathcal{A}_l$ . The final multi-head attention output  $\mathcal{H}_l$  is calculated as a weighted sum of values  $V_l$ . Suppose the two linear layers in FFN are parameterized by  $W_1$ ,  $b_1$  and  $W_2$ ,  $b_2$ , the output of FFN can be formulated as:
$$
\mathcal {F} _ {l} = \operatorname {G e L U} \left(\mathcal {H} _ {l} W _ {1} + b _ {1}\right) W _ {2} + b _ {2} \tag {2}
$$
We term  $\mathcal{F}_l \in \mathbb{R}^{N \times D}$  as the output feature of the  $l$ -th Transformer block, where  $D$  denotes the dimension of hidden features. Some works directly adopted  $\mathcal{F}_l$  as the distillation target. However, the embedding dimensions  $D_S$  and  $D_T$  of student and teacher are typically different. Previous works Romero et al. [2014], Yim et al. [2017], Heo et al. [2019] overcome this obstacle by building certain adaptation modules between hidden layers of the teacher and student models. However, these adaptation modules, with random initialization or special non-parameter transformation Srinivas and Fleuret [2018], Komodakis and Zagoruyko [2017] would potentially disturb training process, because it introduces extra parameters and computational cost (including weights, gradients and optimizer states) Pudipeddi et al. [2020]. Moreover, the teacher and student models usually have different number of heads, i.e.,  $H_T \neq H_S$ . To address these challenges, we aim to model feature relationships using  $\mathcal{F}_l$  to overcome the aforementioned issues. The details of this process are provided in the following sections.
3.2 Distillation with feature relationships
Token-level Relation To mitigate the negative influences of magnitude differences, we first normalize the block features  $\mathcal{F}$  of both the teacher and the student, denoted as  $\hat{\mathcal{F}} = \mathrm{Norm}(\mathcal{F})$ . Common
normalization methods include  $\ell_2$ , softmax and layer normalization. In our implementation, we choose  $\ell_2$  normalization as its implementation, as it consistently outperforms the other methods in our experiments. Subsequently, with the normalized feature, we compute the token-level relation matrix. Specifically, for a specific sample  $i$ , we define the relation matrix between tokens for each sample as follows:
$$
\mathcal {R} _ {t} (\mathcal {F}) = \frac {\mathcal {F} _ {i , ; ;}}{\| \mathcal {F} _ {i , ; ; ;} \| _ {2}} \cdot \left(\frac {\mathcal {F} _ {i , ; ;}}{\| \mathcal {F} _ {i , ; ; ;} \| _ {2}}\right) ^ {T r} \tag {3}
$$
Here,  $Tr$  denotes transposition for the feature. For each given sample, both the student's and teacher's token-level relation matrices share the same dimensions of  $\mathcal{R}_t\in \mathbb{R}^{N\times N}$ . Each matrix effectively serves as a relevance map, revealing the influence of each token in relation to others within the same sequence (see the left part of Figure 2). By doing so, this inner token-level relationship is adept at capturing long-term dependencies between tokens, emphasizing those which are crucial for a comprehensive understanding of the linguistic context. Consider tasks such as text tagging or named entity recognition, where each token must be assigned a label based on its role within the sentence. In such cases, the token-level relation matrix can serve as an effective guide, assisting the student model to deliver enhanced performance.
Sample-level Relation In addition to the token-level relation, which captures the relationships among different tokens within each sample, the relationships across multiple samples for each token also provides informative knowledge for the distillation process. As such, we aim to distill this sample-level relation as well to enhance performance. Similarly, we compute the relation matrix at the sample level as follows:
$$
\mathcal {R} _ {s} (\mathcal {F}) = \frac {\mathcal {F} _ {: , j ;}}{| | \mathcal {F} _ {: , j ;} | | _ {2}} \cdot \left(\frac {\mathcal {F} _ {: , j ;}}{| | \mathcal {F} _ {: , j ;} | | _ {2}}\right) ^ {T r} \tag {4}
$$
In this case, for a specific token, both the student's and teacher's sample-level relation matrices, denoted as  $\mathcal{R}_s$ , share the same dimensions of  $\mathbb{R}^{B\times B}$ . These matrices describe the relationships between different samples, as illustrated in the right part of Figure 2. They enable the capture of semantic relationships across a batch of samples, which is particularly essential for tasks like text classification and summarization. In these tasks, recognizing the sample-level relationship can be instrumental in making accurate predictions. It allows the model to discern the similarities and differences between samples, rather than viewing each sample in isolation. For the sample-level relation modeling, instead of directly comparing words based on their positions across sentences, we transform each sentence into a unified high-dimensional space using the same network. Within this space, token features at the same positions from distinct sequences become comparable. This approach forms the foundation of our sample-level relation modeling and serves as a valuable guide for the student model, enhancing its performance in tasks necessitating a deep understanding of inter-sample relationships.
Computation Complexity Analysis The token-level and sample-level relation maps, in our proposed method, require computational complexities of  $2BN^{2}D$  and  $2B^{2}ND$ , respectively. The associated memory space required for these computations is  $BN^{2} + B^{2}N$ . In contrast, MiniLM Wang et al. [2020b] utilizes self-attention and value-value attention, results in total computational complexities and memory space requirements of  $4BN^{2}D$  and  $2BN^{2}$ , respectively. MiniLMv2 Wang et al. [2020a] employs a different attention mechanism, resulting in total computational complexities and memory space requirements of  $18BN^{2}D$  and  $9BN^{2}$ , respectively. For a concrete example, consider the  $\mathrm{BERT}_{\mathrm{BASE}}$  model, where  $B$  is set to 32,  $N$  is 128, and  $D$  is 768. Our proposed method yields computational complexities and memory space requirements of 0.6GFLOPs and 0.6MB, respectively. In comparison, MiniLM has computational complexities and memory space requirements of 3.2GFLOPs and 2MB, respectively, while for MiniLMv2, these values rise to 14.5GFLOPs and 9.4MB, respectively. Thus, our method demonstrates significant efficiency in terms of computational complexity and memory space requirement, providing a more resource-efficient option for knowledge distillation.
3.3 Distillation with pearson linear correlation
In addition the distillation target, the design of the distillation loss plays a crucial role in transferring the knowledge from the teacher model to the student model. A general distillation loss can be
expressed as:
$$
\mathcal {L} = \Gamma \left(\phi_ {T} \left(\mathcal {F} _ {T}\right), \phi_ {S} \left(\mathcal {F} _ {S}\right)\right) \tag {5}
$$
Here,  $\phi_T$  and  $\phi_S$  denote the feature transformations for the teacher and student models, respectively. These transformations convert raw features into a form that is more conducive to knowledge transfer. In previous works,  $\phi$  often takes the form of an adaptation layer that aligns the embedding dimensions between the teacher and student models. However, in our method, as discussed in Section 3.2,  $\phi$  is used to denote the token-level and sample-level relation matrix.  $\Gamma$  is a distance function measuring the similarity between student and teacher features. The most commonly used functions include KL divergence and the mean squared error. However, these functions exhibit an 'exact match' property, which means the loss reaches the minimal if and only if the features of student and teacher are all identical. This requirement could impose overly strict constraints on the student, particularly when there is a large discrepancy between the teacher and student. To alleviate these constraints, we introduce the Pearson correlation coefficient as a more flexible alternative, which is used to measure the strength of the linear relationship between two variables and remains invariant under positive linear transformations. The basic form of Pearson correlation coefficient between  $X$  and  $Y$  follows the Pearson index can be computed as follows:
$$
\rho (X, Y) := \frac {\sum \left(X _ {i} - \mu_ {X}\right) \left(Y _ {i} - \mu_ {Y}\right)}{\sqrt {\sum \left(X _ {i} - \mu_ {X}\right) ^ {2}} \sqrt {\sum \left(Y _ {i} - \mu_ {Y}\right) ^ {2}}} \tag {6}
$$
where  $\mu_{X},\mu_{Y}$  denote the mean value of the variable  $X,Y$  respectively. The Pearson correlation coefficient, denoted by  $\rho$ , ranges between -1 and 1. Therefore,  $1 - \rho$  always falls between 0 and 2, making it suitable to be used as a loss function. We call our final formulation of the distance function Pearson linear correlation (PLC), which is defined as follows:
$$
\Gamma (X, Y) := 1 - \rho (X, Y). \tag {7}
$$
In this way, we shift the focus from attempting to exactly replicate the teacher's features to preserving and learning the relational information between the teacher's and student's features. This shift effectively relaxes the 'exact match' requirement inherent in conventional Knowledge Distillation (KD) methods. Thus, the distillation loss in our approach comprises two types of losses: token-level relation loss and sample-level relation loss. The token-level relation loss quantifies the discrepancy between the token-level relation matrices of the student and teacher models. It is defined as the average PLC between the student's and teacher's token-level relations across all samples in the batch:
$$
\mathcal {L} _ {t} := \frac {1}{B} \sum_ {i = 1} ^ {B} \Gamma \left(\mathcal {R} _ {t} \left(\mathcal {F} _ {S}\right), \mathcal {R} _ {t} \left(\mathcal {F} _ {T}\right)\right) \tag {8}
$$
The sample-level relation loss quantifies the discrepancy between the sample-level relation matrices of the student and teacher models. It is defined as the average PLC between the student's and teacher's sample-level relations across all tokens in the sequence.
$$
\mathcal {L} _ {s} := \frac {1}{N} \sum_ {j = 1} ^ {N} \Gamma \left(\mathcal {R} _ {s} \left(\mathcal {F} _ {S}\right), \mathcal {R} _ {s} \left(\mathcal {F} _ {T}\right)\right) \tag {9}
$$
The intrinsic sensitivity of the PLC to outliers necessitates normalization of features prior to calculating the distillation loss, thus ensuring a stable training process. Given these two components, the overall training loss of our proposed method consists of the task loss, token-level relation loss, and sample-level relation loss, which can be formulated as follows: