Entity-enhanced BERT for medical specialty prediction based on clinical questionnaire data

Abstract

A medical specialty prediction system for remote diagnosis can reduce the unexpected costs incurred by first-visit patients who visit the wrong hospital department for their symptoms. To develop medical specialty prediction systems, several researchers have explored clinical predictive models using real medical text data. Medical text data include large amounts of information regarding patients, which increases the sequence length. Hence, a few studies have attempted to extract entities from the text as concise features and provide domain-specific knowledge for clinical text classification. However, it is still insufficient to inject them into the model effectively. Thus, we propose Entity-enhanced BERT (E-BERT), which utilizes the structural attributes of BERT for medical specialty prediction. E-BERT has an entity embedding layer and entity-aware attention to inject domain-specific knowledge and focus on relationships between medical-related entities within the sequences. Experimental results on clinical questionnaire data demonstrate the superiority of E-BERT over the other benchmark models, regardless of the input sequence length. Moreover, the visualization results for the effects of entity-aware attention prove that E-BERT effectively incorporate domain-specific knowledge and other information, enabling the capture of contextual information in the text. Finally, the robustness and applicability of the proposed method is explored by applying it to other Pre-trained Language Models. These effective medical specialty predictive model can provide practical information to first-visit patients, resulting in streamlining the diagnostic process and improving the quality of medical consultations.

Fig 1. Overall framework for our proposed method.

The architecture of E-BERT (left), and the detailed implementation of the E-BERT embedding block and entity-aware multi-head self-attention module (right).

Fig 2. Medical NER model for sequence-level prediction.

Fig 3. Comparison of prediction performance according to input sequence length in benchmark models and E-BERT.

The samples were categorized as “short”, “middle”, and “long” based on the quartiles Q1 and Q3 of the sequence length distribution, considering their relative lengths in the dataset.

Fig 4. Visualization of the attention scores averaged across all heads and layers for two random samples in BERT (left) and E-BERT (right), respectively.

In the case of E-BERT, we applied entity-aware attention to 7–9 layers of all the encoder layers. For the mapping between Korean tokens and their English translations, please refer to S1 Table.

Fig 5. Gate values in each layer of E-BERT in which entity-aware attention is applied in all layers.

The tokens of the medical-related entity type have larger gate values than other tokens in the early and middle layers (left). Specifically, the tokens of location type have the largest value in most layers (right).