Keywords

1 Introduction

Healthcare text mining attracts increasing research interest as electronic health record (EHR) and healthcare claim data have skyrocketed over the past decade. Recently, deep pre-trained language models, such as BERT [2] and GPT [3], have improved many natural language processing tasks significantly. However, it won’t give satisfactory results by directly applying those deep pre-trained language models to healthcare text mining. One important reason is that those models are trained from generic domain corpora, which contains a word distribution shift from healthcare corpora. Moreover, deep pre-trained language models are difficult to use on resource-restricted devices due to their huge computation complexity and memory consumption. It’s very important to have embedded models that can directly inference on mobile for healthcare related apps in the US because: 1) it can provide better user experience at poor cell phone signal locations, and 2) it doesn’t require users to upload their health sensitive information onto the cloud. In the US, health related data are only allowed to upload to the cloud by mobile apps being developed by certified institutes, which greatly suppresses the enthusiasm of developing healthcare mobile apps from individual developers. There are some model compression techniques developed recently for generic BERT [6,7,8], but there doesn’t exist a small and efficient enough pre-trained language model in healthcare domain. In this work, we developed SqueezeBioBERT. SqueezeBioBERT has 3 transformer layers, and inference much faster while being accurate on healthcare natural language processing tasks. Our contributions are summarized as below:

  • We designed a novel knowledge distillation method, which is very effective for compressing Transformer-based models without losing accuracy.

  • We applied this knowledge distillation method to BioBERT [5], and experiments show that knowledge encoded in the large BioBERT can be effectively transferred to a compressed version of SqueezeBioBERT.

  • We evaluated SqueezeBioBERT on three healthcare text mining tasks: name entity recognition, relation extraction and question answering. The result shows that SqueezeBioBERT achieves more than 95% of the performance of teacher BioBERT on these three tasks, while being 4.2X smaller.

2 Transformer Layer

As the foundation of modern pre-trained language models [2,3,4], transformer layer [1] can capture long-term dependencies of the input tokens with attention mechanism. A typical transformer layer contains two major components: multi-head attention (MHA) and feed-forward network (FFN).

2.1 Multi-head Attention

Practically, we calculate the attention function on a query set Q, with key set K and value set V. The attention function can be defined as below:

$$\begin{aligned} \mathbf{A} = \frac{\mathbf{Q }\mathbf{K ^{T}}}{\sqrt{d_k}} \end{aligned}$$
(1)
$$\begin{aligned} Attention(\mathbf{Q} ,\mathbf{K} , \mathbf{V} ) = softmax(\mathbf{A} )\mathbf{V} \end{aligned}$$
(2)

where \(d_k\) denotes the dimension of \( \mathbf{K} \).

Multi-head attention will jointly train the model from different representation subspaces. It is denoted as below:

$$\begin{aligned} MultiHead(\mathbf{Q} ,\mathbf{K} , \mathbf{V} ) = Concat(head_1, ..., head_h)\mathbf{W} \end{aligned}$$
(3)

where h denotes attention head number, \(head_i\) is computed by Eq. (2), and W is the linear parameter weight.

2.2 Feed-Forward Network

After multi-head attention, a fully connected feed-forward network will follow, which is denoted as below:

$$\begin{aligned} FFN(x) = max(0, x\mathbf{W} _1 + b_1)\mathbf{W} _2 + b_2 \end{aligned}$$
(4)

3 Knowledge Distillation

A very common way to boost the performance of a machine learning algorithm is to train several models, and then ensemble. Deep learning models are generally heavy neural networks, so it’s normally considered too computationally expensive and inefficient to deploy the ensemble of deep neural networks in the production environment. [9] first proposed Knowledge Distillation and showed the possibility of compressing the function learned from a large complex model into a much smaller and faster model without significant accuracy loss [10]. As deep learning models are becoming more and more complex, knowledge distillation has shown its power of transferring the knowledge from a group of specialist networks to a single model [10,11,12].

Formally, Knowledge Distillation process can be defined as the process of minimizing the loss function between the a large teacher network T and a small student network S as below:

$$\begin{aligned} \mathcal {L}_{KD} = \sum _{x \in X}L(f^T(x), f^S(x)) \end{aligned}$$
(5)

where L denotes the loss function to evaluate the difference between T and S, x is the token input, X is the training set, \(f^T\) denotes the output of the teacher network T and \(f^S\) denotes the output of the student network S.

4 BioBERT

BioBERT [5], with almost the same structure as BERT and pre-trained on biomedical domain corpora such as PubMed Abstracts and PMC full-text articles, can significantly outperform BERT on biomedical text mining tasks.

BioBERT has been fine-tuned on the following three tasks: Named Entity Recognition (NER), Relation Extraction (RE) and Question Answering (QA). NER is to recognize domain-specific nouns in a corpus, and precision, recall and F1 score are used for evaluation on the datasets listed in Table 1. RE is to classify the relationships of named entities, and precision, recall and F1 score are used for evaluation on the datasets listed in Table 2. QA is to answer a specific question in a given text passage, and strict accuracy, lenient accuracy and mean reciprocal rank are used for evaluation on BioASQ factoid dataset [24].

Table 1. BioBERT Named Entity Recognition evaluation datasets
Full size table
Table 2. BioBERT Relation Extraction evaluation datasets
Full size table

5 BioBERT Distillation

In this section, we developed a novel distillation method for BioBERT. Experiments show that knowledge encoded in the large BioBERT can be effectively transferred to the compressed version of SqueezeBioBERT.

Figure 1 shows the overview of the proposed knowledge distillation method. Supposing that the teacher BioBERT has M transformer layers and the student SqueezeBioBERT has N transformer layers, we distillated BioBERT both on transformer layers and task-specific layers.

Fig. 1.
figure 1

The overview of distillation from BioBERT to SqueezeBioBERT

Full size image

Transformer layer distillation consists of multi-head attention distillation and feed forward network distillation. For multi-head attention distillation, we combine Eqs. (2), (3) and (5), and use the mean squared error (MSE) as the loss function since it’s more suitable for regression tasks. Thus, the multi-head attention distillation process is denoted as below:

$$\begin{aligned} \mathcal {L}_{MHA} = \frac{1}{h} \sum _{i = 1}^{h}MSE(M^T_i,M^S_i) \end{aligned}$$
(6)

where h denotes the number of attention heads, \(M^S_i\) denotes the output of i-th student attention head, and \(M^T_i\) denotes the output of i-th teacher attention head.

For feed forward network distillation, we can use a single linear transformation \(W_{FFN}\) to transform the output of the teacher network into the student network. Thus, the feed forward network distillation process is denoted as below:

$$\begin{aligned} \mathcal {L}_{FFN} = MSE(O^T_{MHA}W_{FFN},O^S_{MHA}) \end{aligned}$$
(7)

For task-specific prediction layer distillation, we use softmax cross-entropy as the loss function, since it’s more suitable for classification tasks. Thus, the task-specific prediction layer distillation is denoted as below:

$$\begin{aligned} \mathcal {L}_{pred} = -softmax(O^T_{FFN})log(softmax(O^S_{FFN})) \end{aligned}$$
(8)

In summary, Eqs. (6), (7) and (8) describes the overall procedure of the BioBERT distillation process.

Table 3. Named Entity Recognition metrics comparison
Full size table
Table 4. Relation extraction metrics comparison
Full size table

6 Experiments

We use BioBERT-Base v1.1 [25] as our source model, and distillated it to SqueezeBioBERT on the same three healthcare NLP tasks. BioBERT-Base v1.1 has 12 transformer layers and 109M weights. SqueezeBioBERT has 3 transformer layers and 26M weights.

Table 5. Question answering metrics comparison
Full size table

NER results are show in Table 3, RE results are show in Table 4, and QA results are show in Table 5. From the results, we can see that SqueezeBioBERT is 4.2X smaller than BioBERT, but still achieves more than 95% accuracy performance of the teacher BioBERT on the three NLP tasks. This proves the efficiency of the proposed method of transferring knowledge encoded in the large BioBERT to the compressed version of SqueezeBioBERT.

7 Conclusion

Although recent deep pre-trained language models have greatly improved many natural language processing tasks, they are generally computationally expensive and memory intensive, which makes them very difficult to use on resource-restricted mobile or IoT devices. Embedded models that can directly inference on mobile is important for healthcare related apps in the US because: 1) it can provide better user experience at poor cell phone signal locations, and 2) it doesn’t require users to upload their health sensitive information onto the cloud. In this paper, we designed a novel knowledge distillation method, which is very effective for compressing Transformer-based models without losing accuracy. We applied this knowledge distillation method to BioBERT, and experiments show that knowledge encoded in the large BioBERT can be effectively transferred to a compressed version of SqueezeBioBERT. We evaluated SqueezeBioBERT on three healthcare text mining tasks: name entity recognition, relation extraction and question answering. The result shows that SqueezeBioBERT achieves more than 95% of the performance of teacher BioBERT on these three tasks, while being 4.2X smaller.