Towards Non-task-specific Distillation of BERT via Sentence
Representation Approximation

Suppose $x=\{w_{1},w_{2},\cdots,w_{i},\cdots w_{n}|i\in[1,n]\}$ stands for a sentence containing $n$ tokens (${w_{i}}$ is the $i$-th token of $x$), and let $T:x\to T_{x}\in\mathbb{R}^{d}$ be the teacher model which encodes $x$ into $d$-dimensional sentence embedding $T_{x}$, the goal of the sentence approximation oriented distillation is to train a student model $S:x\to S_{x}\in\mathbb{R}^{d}$ generating $S_{x}$ as the approximation of $T_{x}$.

In our proposed distillation architecture, as shown in Figure 1a, we take the BERT as the teacher model $T$, and the hidden representation $C$ is extracted from the top transformer layer upon the [CLS]¹¹1[CLS] is a special symbol added in front of other tokens in BERT, and the final hidden state corresponding to this token is usually used as the aggregate sequence representation. token as $T_{x}$. For the student model, a standard bidirectional LSTM (BiLSTM) is first employed to encode the sentence into a fixed-size vector $H$. After that, a fully connected layer without bias terms is built upon the BiLSTM layer to map $H$ into a $d$-dimensional representation, followed by a $tanh$ activation which normalizes the values in previous representation between -1 and 1 as the final $S_{x}$.

As our non-task-specific distillation task has no labeling data, and the signal given by the teacher is a real value vector, it is not feasible to minimize the cross-entropy loss over the soft labels and ground truth labels Sun et al. (2019); Barkan et al. (2019); Tang et al. (2019b). On this basis, we propose an adjusted cosine similarity between the two real value vectors $T_{x}$ and $S_{x}$ to perform the sentence representation approximation. Our distillation objective is computed as follows:

Here $tanh$ is chosen as the activation function since most values (more than 98% according to our statics) in $T_{x}$ obtained from BERT are within range of $tanh$ (-1 to 1). The choice of using cosine similarity based loss is mainly based on the following two considerations. Firstly, since 2% values in $T_{x}$ are outside the range of [-1, 1], it is more reasonable to use a scalable measurement, such as cosine similarity, to deal with these deviations. Secondly, it is meaningful to compute the cosine similarity between sentence embeddings given by BERT Xiao (2018).

Overall, after the distillation procedure, we obtained a BiLSTM based “BERT”, which is smaller in parameter scale and more efficient in generating a semantic representation for a sentence.

Distilling data

As our distillation procedure needs no dependency on sentence type or labeling resources but only standard sentences available everywhere, the distillation data selection follows the existing literature on language model pre-training as well as BERT. We use the English Wikipedia to perform the distillation. Furthermore, as the proposed method focus on the sentence representation approximation, the document is segmented into sentences using spacy Honnibal and Montani (2017).

The fine-tuning on sentence-level tasks is straightforward. The downstream tasks discussed in this paper can be summarized as two types: type judgment on a single sentence and predicting the relationship between two sentences (same as all the tasks of GLUE). Figure 1b illustrates the model architecture for single sentence classification tasks. The student model $S$ is utilized to provide sentence representation. After that, a multilayer perceptron (MLP) based classifier using Relu as activation of hidden layers is applied for the specific task. For the sentence pair tasks, as shown in Figure 1c, the representations $H$ and $\tilde{H}$ for the sentence pair are obtained by transforming two sentences into two BiLSTM based student models with shared weights respectively. Then, following the baseline BiLSTM model reported by GLUE Wang et al. (2019), we apply a standard concatenate-compare operation between the two sentence embeddings and get an interactive vector as $[H,\tilde{H},|H-\tilde{H}|,H\odot\tilde{H}]$, where the $\odot$ demotes for the element-wise multiplication. Then, same as the single sentence task, an MLP based classifier is built upon the interactive representation.

For both types of tasks, MLP layers are initialized randomly, and the rest parameters are inherited from the distilled student model. Meanwhile, all parameters are optimized through the training procedure for the specific task.

To evaluate the performance of our proposed non-task-specific distilling method, we conduct experiments on three types of tasks: sentiment classification (SST-2), similarity (QQP, MRPC), and natural language inference (MNLI). All the tasks come from the GLUE benchmark Wang et al. (2019).

SST-2

Based on the Stanford Sentiment Treebank dataset Socher et al. (2013), the SST-2 task is to predict the binary sentiment of a given single sentence. The dataset contains 64k sentences for training and remains 1k for testing.

QQP

The Quora Question Pairs²²2https://www.quora.com/q/quoradata/First-Quora-Dataset-Release-Question-Pairs dataset consists of pairs of questions, and the corresponding task is to determine whether each pair is semantically equivalent.

MNLI

The Multi-Genre Language Inference Corpus Williams et al. (2018) is a crowdsourced collection of sentence pairs with textual entailment annotations. There are two sections of the test dataset: matched (in-domain, noted as MNLI-m) and mismatched (cross-domain, noted as MNLI-mm).

MRPC

The Microsoft Research Paraphrase Corpus Dolan and Brockett (2005) is similar to the QQP dataset. This dataset consists of sentence pairs with binary labels denoting their semantic equivalence.

BERT Devlin et al. (2019)

with two variants: BERT${}_{\text{BASE}}$ and BERT${}_{\text{LARGE}}$, containing 12 and 24 layers of Transformer respectively.

ELMO Baseline Wang et al. (2019)

is a BiLSTM based model, taking ELMo Peters et al. (2018) embeddings in place of word embeddings.

BERT-PKD Sun et al. (2019)

proposes a patient knowledge distillation approach to compress BERT into a BERT with fewer layers. BERT₃-PKD and BERT₆-PKD stand for the student models consisting of 3 and 6 layers of Transformer, respectively.

DSE Barkan et al. (2019)

is a sentence embedding model based on knowledge distillation from cross-attentive models. For each single sentence modeling, the 24-layers BERT is employed.

BiLSTM${}_{\text{KD}}$ Tang et al. (2019b)

introduces a new distillation objective to distill a BiLSTM based model from BERT for a specific task. BiLSTM${}_{\text{KD}}$+TS Tang et al. (2019a) donates the distilling procedure performed with the proposed data augmentation strategies.

BiLSTM${}_{\text{SRA}}$

stands for the Sentence Representation Approximation based distillation model proposed in this paper. BiLSTM${}_{\text{SRA + KD}}$ donates performing knowledge distillation method proposed by Tang et al. (2019b) during fine-tuning on a specific task, and BiLSTM${}_{\text{SRA + KD}}$+TS demonstrates using the same augmented dataset to perform the distillation.

For the student model in our proposed distilling method, we employ the 300-dimension GloVe (840B Common Crawl version; Pennington et al., 2014) to initialize the word embeddings. The number of hidden units for the bi-directional LSTM is set to 512, and the size of the task-specific layers is set to 256. All the models are optimized using Adam Kingma and Ba (2015). In the distilling procedure, we choose the learning rate as $1\times 10^{-3}$ with the batch size=1024. During fine-tuning, the best learning rate on the validation set is picked from $\{2,3,5,10\}\times 10^{-4}$. For the data augmentation, we use the rule-based method originally suggested by Tang et al. (2019b). Notably, on the SST-2 and MRPC dataset, we stop data augmenting when the transfer set achieves 800K samples following the setting of their follow-up research Tang et al. (2019a). Besides, inspired by the comparisons in the research of Sun et al. (2019), we find BERT${}_{\text{BASE}}$ can provide more instructive representations than BERT${}_{\text{LARGE}}$. So that, we chose BERT${}_{\text{BASE}}$ as our teacher model to pre-train the BiLSTM${}_{\text{SRA}}$.

For a comprehensive experiment analysis, we collect data and implement comparative experiments on various published BERT and BERT-distillation methods. Table 4 shows the results of our proposed BiLSTM${}_{\text{SRA}}$ and the baselines on the four datasets. All models in the first block (row 1-6) belong to base methods without implementing distillation, the second (row 7-9) and third (row 10-12) blocks show the performances of distillation models using BERT and BiLSTM structures respectively, while the forth block (row 13-14) displays the influences of textual data augmentation approach on our BiLSTM${}_{\text{SRA}}$ and BiLSTM${}_{\text{KD}}$ distillation baseline. The last two blocks contain the results of pure improvements obtained by different distillation methods. To analyse the effectiveness of BiLSTM${}_{\text{SRA}}$ thoroughly, we break down the analyses into the following two perspectives.

To compare the inference speeds of different models, we also implement experiments on 100k samples from the QQP dataset, the results are shown in Table 2. The inference procedure is performed on a single P40 GPU with a batch size of 1024. As the inference time is affected by the computing power of the test machine, for fair comparisons with ELMO, BERT₃-PKD, BERT₆-PKD and DSE, we inherit the speed-up factors from previous papers. Besides, the numbers of parameters reported in Table 2 exclude those from the embedding layers, since such components do not affect the inference speed and are highly related to the vocabulary sizes, i.e., usually few words appeared for a specific task.

From the results shown in Table 2, it can be observed that the BiLSTM based distilled models have fewer parameters than BERT, ELMO, as well as the other transformer-based models. Both the BERT${}_{\text{BASE}}$ and ELMO are around 20 times larger in parameter size and 20 times slower in inference speed. Even the smallest transformer based model BERT₃-PKD is also four times larger than our proposed BiLSTM${}_{\text{SRA}}$. Comparing with BiLSTM${}_{\text{KD}}$, although our proposed BiLSTM${}_{\text{SRA}}$ is larger in parameter size due to the restriction of the sentence embedding’s dimension given by the teacher BERT, it stills inferences more efficiently. This is mainly due to the fact that the more hidden units in BiLSTM${}_{\text{SRA}}$ are more accessible to calculated in parallel by the GPU core, while the larger word embedding size in BiLSTM${}_{\text{KD}}$ slows down its inference efficiency. In conclusion, the cost and production per second of BiLSTM${}_{\text{KD}}$ and BiLSTM${}_{\text{SRA}}$ are within the same scale, but our method achieves better results on GLUE tasks according to the comparison shown in Table 4.

Since pre-trained language models have well-initialized parameters and only learn a few parameters from scratch, these models usually converge faster and are less dependent on large-scale annotations. Correspondingly, the non-task-specific distillation method proposed in this paper also aims at obtaining a compressed pre-trained BERT as well as keeping these desirable properties. To evaluate it, in this section, we discuss the influence of the task-specific training data and learning iterations on the performance of our model and the others.

As illustrated in Table 3, we experiment in training the models using different proportions of the dataset. BERT${}_{\text{LARGE}}$ trained on the corresponding data stands for the teacher model of each BiLSTM${}_{\text{KD}}$. With no doubt, all the models can achieve better results using more training data, while BERT performs the best. BERT even successfully predicts 91.9% of validation samples under only 20% training data. Comparing with the pure BiLSTM models, the BiLSTM${}_{\text{KD}}$ models slightly improve the performances by 1%$\sim$2%, whereas BiLSTM${}_{\text{SRA}}$ outperforms the best BiLSTM model as well as the BiLSTM${}_{\text{KD}}$ trained with 20% and 30% percent data respectively. Besides, similar to BERT, the difference of accuracy between BiLSTM${}_{\text{SRA}}$ trained with 20% and the one using 100% corpus is relatively small. This phenomenon indicates that our model converges faster and is less dependent on the amount of training data for downstream tasks.

Such conclusions are also reflected in the comparison in Figure 2 of the models’ learning curves on QQP. Even though QQP is a large dataset to train a good BiLSTM model, it can be observed that BiLSTM${}_{\text{SRA}}$ trained with 30% data performs equivalent to BiLSTM using the whole corpus. Moreover, using 100% training data, BiLSTM${}_{\text{SRA}}$ even outperforms the converged BiLSTM after the first epoch. Besides, all the BiLSTM${}_{\text{SRA}}$ models converge in much fewer epochs.

Despite the task-specified data, Wikipedia corpus is used in the distillation procedure of our proposed method. We also train different BiLSTM${}_{\text{SRA}}$ base models using {1, 2, 4} million Wikipedia data, and the corresponding fine-tuning performances on SST-2 and MNLI are reported in Table 4. It can be observed that both the performances of BiLSTM${}_{\text{SRA}}$ on SST-2 and MNLI are proportional to the distillation loss. This observation indicates the effectiveness of our proposed distillation process and objective. Besides, distilling with adequate data is sufficient to produce more BERT-like sentence representations. Nevertheless, different from the fact that more training data has a significant benefit in a particular task, four times the distilling data only improve around 0.5 points on both two downstream tasks. Thus, our method does not require a vast amount of training data and a long training time to obtain good sentence representations. Furthermore, the loss scores in the second column suggest BiLSTM${}_{\text{SRA}}$ can generate more than 95% similar sentence embeddings with the ones given by BERT under the measure of the cosine distance.

A notable characteristic of the pre-trained language models, such as ELMO, BERT, and certainly the non-task oriented distillation models, lies in the capability of providing sentence representations for quantifying similarities of sentences, without any tuning operation based on specific tasks. In this subsection, we conduct the comparisons among models by directly extracting their sentence embeddings without fine-tuning upon sentence similarity oriented tasks.

Table 5 lists the results of models on the QQP dataset. It should be noted that, in this table, ELMO, BERT${}_{\text{BASE}}$ (CLS) and BERT${}_{\text{BASE}}$ (averaged) are introduced as the comparison basis, since they can give the SOTA untuned sentence representations for the similarity measurement. The comparison mainly focuses on the performances of our proposed BiLSTM${}_{\text{SRA}}$ and BiLSTM${}_{\text{KD}}$. For a thorough comparison, we define the training objective of BiLSTM${}_{\text{KD}}$ as fitting the cosine similarity score of the sentence pair directly given by the pre-trained BERT, which means both the teacher BERT and KD do not utilize the labels of QQP dataset. It can be seen that, even though the training goal of BiLSTM${}_{\text{KD}}$ is more direct than BiLSTM${}_{\text{SRA}}$, our BiLSTM${}_{\text{SRA}}$ outperforms the former on both the metrics, and meanwhile achieves scores closed to those of BERT${}_{\text{BASE}}$. Besides, we can also observe that, for sentence similarity quantification, averaging the context word embeddings as the sentence representation (ELMO and BERT${}_{\text{BASE}}$ (averaged)) works better than taking the final hidden state corresponding to the [CLS] token (BERT${}_{\text{BASE}}$ (CLS)).