A self-attention based neural architecture for Chinese medical named entity recognition

1.   Introduction

With the combination of the medical field and big data, more and more consultation data and disease information are recorded in the form of electronic medical records (EMRs), and gradually become an important basis for assisting doctors in therapeutic diagnosis. EMRs record a large number of diagnostic information of patients: hospital records, course records, doctor's orders, case data and so on, including key entity information such as disease, surgery, drugs, etc. This information is a decisive factor for doctors to make treatment plans for patients ^[1]. It is of great significance to study how to extract key entity information from massive EMRs efficiently and accurately through intelligent methods.

Named entity recognition (NER) is a vital part of natural language processing (NLP) that meets the aforementioned requirements ^[2]. Its purpose is to recognize various named entities, e.g. names, place, organizations, etc., from raw text. Extracted entities can be taken as information for people, and can also pave the way for other NLP tasks, such as relationship extraction and knowledge graph construction. Recently, with the rise of deep learning technology, deep neural networks are utilized to achieve medical NER and have attracted much research attention.

So far, NER still faces huge problems in the field of Chinese electronic medical records (CEMRs). The main reasons are as follows: first of all, an entity may have multiple names due to the undefined text labeling standards ^[3]; secondly, the meaning of the same word or character may be completely different in different contexts, which causes confusions towards Chinese semantics; last but not least, Chinese has no natural vocabulary boundaries (spaces) as English does, so Chinese have no strict and correct vocabulary boundaries. In previous NER researches, the BiLSTM-CRF, which is the abbreviation of bi-directional Long-Short Term Memory (LSTM) joining with a conditional random field (CRF) layer, shows advanced performance and has become a prevalent architecture for various NER tasks ^[4,5]. This architecture outperforms traditional methods in that it eliminates the inefficient and complex method of manually designing feature templates and utilizes recurrent neural network (RNN) to automatically capture text features. However, lengths of CEMR texts are generally longer than traditional text, for CEMR texts contain at least several hundred Chinese tokens. Although it can capture long-term contextual dependencies ^[6], Long-Short Term Memory (LSTM) present inferior performance when text lengths exceed a certain step size ^[7]. In addition, at the beginning of the NLP task, each token in the text is represented by a low-dimensional dense vector ^[8]. Due to the highly domain-specialized medical field, universal pre-trained models are hardly adopted in practical tasks. The lack of Chinese medical corpus even makes it more difficult to pre-training medical domain-specific language models. In most cases, vector representations are generally randomly initialized, leading to a context-independent representation for each token ^[9], as a result, it can't tackle polysemy problems, and the limitations are very obvious.

Our work focus on Chinese medical NER in CEMRs, which has been a subtask of numerous influential academic conferences in NLP domain, e.g. China Conference on Knowledge Graph and Semantic Computing (CCKS), China Health Information Processing Conference (CHIP) etc.. These tasks not only accelerate Chinese medical NER research, but also provide several precious corpora for Chinese medical NER. In this paper, we first collect 4679 CEMRs, which are utilized to fine-tune the Chinese Embeddings from Language Models (ELMo) ^[9] trained with common field corpora, and then get a pre-trained model that can be used to dynamically generate context-dependent character embeddings for Chinese characters. Secondly, encoder from Transformer (ET) ^[10] is utilized as the model encoder instead of the traditional bi-directional Long-Short Term Memory (BiLSTM) in order to provide the proposed model with the ability to capture the long-term dependence of ultra-long CEMR texts efficiently. In ET, the distance between token and token is one, so there is no problem that the dependence is lost due to the lengthy distance between tokens. Our contributions are summarized as follows.

1) We fine-tuned a Chinese medical domain-specific ELMo model, which provides an authentic pre-trained language model for further research. A Chinese medical corpus with 4679 real-world CEMRs is constructed, containing about 1.8 million Chinese characters. Then a medical domain-specific model ELMo is fine-tuned by the efficiently application of Chinese medical corpus and the public available Chinese ELMo model.

2) We realize Chinese medical NER in CEMRs with ET-CRF model, which can tackle the long context dependencies better than the BiLSTM model with self-attention mechanism, and to the best of our knowledge this is the first time to apply ET-CRF model to Chinese medical NER.

3) Owning to the contributions above, the proposed ELMo-ET-CRF model achieves the best performance among all model architectures mentioned in this paper on the CCKS 2019 datasets, and the final F1-score is competitive to the current state-of-the-art performance.

2.   Related word

NER has become one of the important tasks of information retrieval, data mining and NLP owing to of its extraordinary significance ^[11], and various solutions have been proposed in the existing literature.

2.1.   Rule based approach

Matching entities through handwritten rules is the main method to deal with NER tasks in the early stage ^[12]. However, the construction of rules requires a certain level of expertise, and even a domain expert cannot enumerate rules that can model all entities. In addition, rules cannot be migrated because they rely on datasets. Thus, a same set of rules may not work under different datasets. This kind of handcrafted approach always leads to a relatively high system engineering cost.

2.2.   Statistical machine learning based approach

The statistical machine learning method treats NER as a sequence labeling problem by inputting a set of sequences and outputing a set of predicted optimal tags sequences. Traditional methods include hidden Markov models ^[13,14], maximum entropy Markov models ^[15], conditional random fields ^[16,17], and support vector machines ^[18]. The most common implementation is feature template with CRF and different feature templates can be combined to form a new feature template. However, this statistical machine learning based method relies heavily on hand-crafted features, which cost a lot of overhead when find the most appropriate features.

2.3.   Neural network based approach

In recent years, with the increase of computing power, training of deep neural network has become simple and feasible. The propose of word embeddings (e.g. word2vec, glove) make the useage of deep neural networks to deal with NLP tasks become a research focus ^[8,19]. Different from the traditional statistical machine learning training method, the training process of the neural network is regarded as an end-to-end process, which can automatically learn data features and avoid the extra overhead such as feature engineering.

More recently, the model structure which is using bidirectional LSTM to encode data with CRF as the decoder has got the most advanced results in the medical NER task ^[4,5], and the effect is significantly better than the traditional statistical model. The theory of model architecture was first proposed by Collobert et al. ^[20], Huang ^[21] and Lample ^[22] used LSTM-CRF for the first time to deal with sentence-level annotation problems. Ma et al. ^[23] used LSTM-CRF structure for the first time in English NER task, and achieved promising results, and then Dong et al. ^[24] first used LSTM-CRF to handle the Chinese NER task. LSTM has a cell structure and gate mechanism that allows the model to effectively capture long-term dependencies and has certain forgetting capabilities ^[6], and the advantage of splicing CRF after encoder layer is that it can use information that has already occurred during the sequence generation process to ensure that the output value is a set of optimal solution sequences. In addition, Zhang et al. ^[28] investigate a lattice-structured LSTM model for Chinese NER, which encodes a sequence of input characters as well as all potential words that match a lexicon. Zhang et al. ^[29] propose a convolutional attention layer to extract the implicit local context features from character sequence. Liu et al. ^[30] propose a Global Context enhanced Deep Transition architecture for sequence labeling. Qiu et al. ^[3] proposed the RD-CNN-CRF model, which effectively reduced the time required for training without losing the accuracy of the model. In terms of the inconsistency of the label, Ji et al. ^[1] proposed a hybrid model based on the attention mechanism, the attention mechanism effectively alleviates the problem of model accuracy decline caused by label inconsistency.

3.   Methods

In our method, raw CEMRs are used as the input of ELMo to get the vector representation of sentences, sentence features are then extracted by ET before decoded through CRF to generate the annotation sequence. Figure 1 shows the structure of the model, which is detailed in this section.

3.1.   Embeddings from Language Models (ELMo)

In early NLP tasks, the input of RNN is a set of word embeddings generated by Word2Vec ^[8], Glove ^[19] etc. However, these embeddings are context-independent, i.e., each word corresponds to a unique static vector without changing with the context, so these methods are limited in the case of polysemy. The emergence of ELMo effectively solves this problem by applying stacked BiLSTM to model the entire sentence from two directions and mapping the sentences into a sequence of vectors. Since LSTM can capture context dependencies, it ensures that the output embedding sequence has a front-to-back correlation.

Given $S$ sequence of $N$ tokens, $\{{w}_{1}, {w}_{2}, ..., {w}_{N}\}$, these tokens first go through a token layer, which transfers the dimension of original character embedding to the input dimension of BiLSTM according to a weight matrix, then the output of token layer will be sent to a stacked BiLSTM to build a language model, as shown in Figure 2. Suppose the sentence length is $N$, and $L$ represents the number of layers of BiLSTM, at each position t, each layer $l$ of LSTM output a context-dependent hidden vector ${h}_{t, l}^{LM}\left(\left[\overrightarrow{{h}_{t, l}^{LM}}, \overleftarrow{{h}_{t, l}^{LM}}\right]\right)$ where $t = 1, 2, ..., N$and $l = 1, 2, ..., L$. The hidden vector ${h}_{t, l}^{LM}$ of the last layer of LSTM output is used by the softmax layer to predict the token at the next moment.

When the training process finishes, each sentence will learn$2L+1$ representations ${R}_{t}$, where $L$ is the number of layers in the model:

where ${h}_{t, 0}^{LM}$ is the output of token layer and ${h}_{t, l}^{LM} = \left[\overrightarrow{{h}_{t, l}^{LM}}, \overleftarrow{{h}_{t, l}^{LM}}\right]$ for each BiLSTM layer.

3.2.   Encoder from Transformer (ET)

Our model uses ET as the encoder layer. Compared with RNN and CNN, it takes a different approach that relies entirely on the self-attention mechanism to extract context features, which makes it highly parallel in the encoding process.

Suppose the input $S$ is a set of sequences $\{{w}_{1}, {w}_{2}, ..., {w}_{N}\}$, $S\in {\mathbb{R}}^{N\times {d}_{model}}$, where N is the length of the sequence and ${d}_{model}$ is the dimension of the input vector. We use multiple of the scaled dot-product attention components inside the multi-head attention layer to enhance the model's ability to encode sequences internally. We first add position information to the input sequence S, which follows the approach of Vaswani et al. ^[10], then perform a matrix transformation operation on S to obtain three weight matrices, namely key matrix $K$, query matrix $Q$ and value matrix $V$. Finally, the output representation of a single self-attention is obtained by scaled dot-product attention:

where ${W}^{Q}\in {\mathbb{R}}^{{d}_{model}\times {d}_{k}}$, ${W}^{K}\in {\mathbb{R}}^{{d}_{model}\times {d}_{k}}$, ${W}^{V}\in {\mathbb{R}}^{{d}_{model}\times {d}_{k}}$, are learnable parameters and $softmax\left(\right)$is performed row-wise.

Single head attention may inhibit the information from different representation subspaces at different positions, so our model uses multi-head attention mechanisms ^[10]:

where ${head}_{i} = Attn({Q}_{i}, {K}_{i}, {V}_{i})$ and ${W}^{O}$ is learnable parameter.

Multi-head attention will be followed by a feedforward neural network, and the output of each sublayer in the encoder will have a residual link and layer normalization as shown in Figure 3.

Assuming that the output of the multi-attention is the $\stackrel{-}{S}$, the final output of the encoder will be calculated as follows:

where $FFN\left(x\right) = \mathrm{max}\left(0, x{W}_{1}+{b}_{1}\right){W}_{2}+{b}_{2}$ and $layernorm\left(\right)$ represents the layer normalization ^[25].

3.3.   Conditional random field (CRF)

In the NER task of CEMRs, the output tag sequence is strictly ordered. The CRF layer maintains a state transition matrix which stores the transition probability from the previous state to the current state, ensureing the tag prediction process has inner dependency.

Given the input sequence is $X = \{{x}_{1}, {x}_{2}, ..., {x}_{N}\}$, the score for defining the output sequence $Y = \{{y}_{1}, {y}_{2}, ..., {y}_{N}\}$ is expressed as follows :

where $A$ is a matrix of transition scores, ${A}_{ij}$ represents the score of a transition from the tag $i$ to tag $j$. $P$ is a matrix of tag scores which is obtained from the output of ET through a fully connected network, $P\in {\mathbb{R}}^{N\times k}$, where $N$ is the length of the sequence, $k$ is the number of labels, and ${p}_{ij}$represents the score of the ${j}^{th}$ label of the ${i}^{th}$position in the sequence. ${y}_{0}$ and ${y}_{N+1}$ are represented by < bos > and < eos > .

The model uses softmax to calculate the probability of all the tag sequences that may be generated by the input sequence X, and defines the logarithmic probability of maximizing the correct annotated sequence as the goal of the model optimization ^[22]:

where ${Y}_{ALL}$ represents all possible tag sequences for an input sentence $X$.

4.   Experiments

4.1.   Fine-tuning ELMo

The Chinese ELMo model used in this paper comes from the Research Center for Social Computing and Information Retrieval Harbin Institute of Technology. It was trained on Xinhua proportion of Chinese gigawords-v5, and takes roughly 3 days on an NVIDIA P100 GPU ^[26,27]. Since the datasets used in this paper belong to medical field, the corpus distribution has a strong specialized background, and there are also some uncommon characters in the text. Therefore, the pre-trained ELMo which is used in the general field cannot be directly applied to this task. We fine-tuned the ELMo with the medical corpus.

We collected 4679 CEMRs, which contains about 1.8 million characters, from CCKS2018, CCKS2019 and CHIP2018, and routinely preprocessed the text. In particular, it should be noted that in this paper the proposed neural network model is based on character embeddings. In the fine-tuning process, the hyperparameter setting is shown in the Table 1.

After the model training is completed, we freeze the parameters of ELMo and use the weighted summation of each layer of ELMo as the output vector of ELMo, finally take the concatenation of the ELMo output and the original character embeddings as downstream model input:

where $\theta$ are softmax-normalized weights and the $\lambda$ is a scalar which can scale the vector according to a certain ratio. These two parameters are learnable and is updated with downstream model training. ${c}_{t}$ represents original character embedding which generated by Word2Vec, ${c}_{t}^{embed}$ represents the final vector representation of the character at the t position in a sequence and it will be the input of ET.

4.2.   Task and dataset

This paper focus on Chinese medical NER, which aims to detect entity boundary and categorize entity into pre-defined categories. The dataset used in this paper comes from CCKS2019, which is jointly provided by Yiducloud (Beijing) Technology Co., Ltd.

The training data contains 1000 manually annotated CMERs, while the test data contains 379 manually annotated CMERs. Each EMR contains two parts: raw CEMR and annotation information. The annotation information consists of several triples, which are formed of entity start index, entity end index and entity category. Through entity start and end indices, we can extract the entity from CMERs. All entities are categorized into six categories, disease and diagnosis, imaging examination, laboratory test, surgery, medicine and anatomy. Entities contained in this dataset is shown in Table 2.

4.3.   Evaluation criteria

We use the standard evaluation criteria to validate the effectiveness of the model, namely precision, recall and micro F1-score. which can be calculated as follows:

where the $TP, TN, FP$and $FN$are true positive, true negative, false positive and false negative respectively.

In this paper, CMERs are encoded with BIO format (Begin, Inside and Outside). The token will label as B-label if the token is the beginning of a named entity, I-label means the token is inside a named entity, other case will label as O-label. It can be seen from Figure 4 that tag ‘O’ are negative samples, and other tags are positive samples.

5.   Results and discussions

We first use the training data to train the BiLSTM-CRF. In the experiment, the batch size was 10, Adam was selected as the model's optimizer, and the learning rate was set to 0.001. The initial character embedding is generated by word2vec method with a dimension of 300. The BiLSTM layer number is 1, and the hidden layer vector dimension is 300.

For ET-CRF, since the layer numbers and self-attention heads have a critical impact on the model's performance, we trained six sets of models with different numbers of layers or heads, and compared the effects of the models under different parameters on the test set. The results are shown in Table 3. The hyperparameters for the best model are 2 layers and 8 heads. In addition, the model’s input is 512 dimensions character embedding, which is also generated by the word2vec method with the same experimental configuration as BiLSTM-CRF.

The results of WV-BiLSTM-CRF and WV-ET-CRF are shown in Table 4 (WV represents word2vec). It is obvious that the F1-score of the medicine and the imaging examination part is ideal for the BiLSTM-CRF, but the recognition ability of disease and diagnosis is the short board of the model, F1-score only 76.23%. It is speculated that the entity length of the medicine and the anatomy part is generally short and there is no multiple standard, but the tag of the entity such as disease and diagnosis is subjective, context-dependent, and generally has a long length, so judging the boundary is quite difficultly for BiLSTM. The total F1-score of the ET-CRF is 84.05%, which is higher than that of BiLSTM-CRF, 2.19%, revealing that about 382 medical entities are rectified or extracted. And for entities such as disease and diagnosis, ET-CRF performs remarkable. This is an intuitive result, we speculated that due to the ET encoding process which offers effectively method of shortening the direct distance between characters, the context-dependent relationship become well captured and preserved, ensuring the good recognition rate of ET-CRF for longer entities.

In order to verify the actual effect of ELMo on the NER task in CEMRs, we add ELMo components to the above two models, BiLSTM-CRF and ET-CRF. As a result, the input of the model becomes a dynamic context-dependent character embeddings with context information. The test results are shown in Table 5. For the ELMo-BiLSTM-CRF, the total F1-score is 2.64% higher than BiLSTM-CRF, which means that about 460 medical entities are rectified or extracted. The addition of ELMo has also improved the recognition ability of disease and diagnosis, which can be seen from the table that 2.24% has been improved. For the ELMo-ET-CRF, the total F1-score is 1.54% higher than that of ET-CRF, which shows that 268 medical entities are rectified or extracted. In summary, the addition of pre-trained ELMo enriches the information contained in the character embedding, which effectively improves the accuracy of the model.

Table 6 illustrates that ELMo outperforms word2vec due to the obtained dynamic context-dependent model input,and the performance of ET-CRF in Chinese medical named entity recognition is significantly better than that of BiLSTM-CRF. Due to the excellent long context dependency capture capability of the self-attention mechanism,the ET-CRF’s ability to recognize long entity is significantly better than the BiLSTM-CRF. In addition,we also found that the convergence speed of ET-CRF is significantly faster than BiLSTM-CRF. It can also be seen in the table that the final F1-score of the best model ELMo-ET-CRF in the paper is 85.59%,which is still competitive compared to the top three in the CCKS2019 competition (https://www.biendata.com/competition/ccks_2019_1/final-leaderboard/).

6.   Conclusions

In this paper, we firstly fine-tune a medical domain-specific ELMo model through a small medical corpus which is contained with 4679 CEMRs. Then we apply the ET-CRF model to Chinese medical NER on CEMRs. Finally, the proposed ELMo-ET-CRF model use dynamic context-dependent ELMo character embeddings to incorporate more lexical, syntantic and semantic information, and alleviates long context dependency problem. Under the strict evaluation index, the F1-score of ELMo-ET-CRF on the test set is 85.59%, which is competitive to the state-of-the-art on this dataset, and indicates the effectiveness of the proposed model architecture.

Acknowledgments

We thank the anonymous reviewers for their constructive comments. This work was supported by grants from the National Key Research and Development Program of China (2017YFB0202104). Publication costs are funded by a grant of the National Key Research and Development Program of China (2017YFB0202104).

Conflict of interest

All authors declare no conflicts of interest in this paper.