Preserving Empirical Probabilities in BERT for Small-sample Clinical Entity Recognition
Abdul Rehman Jian Jun Zhang Xiaosong Yang

Gr-IoU: Ground-Intersection over Union for Robust Multi-Object Tracking with 3D Geometric Constraints
    Keisuke Toida10009-0006-4873-3651 Naoki Kato20009-0004-3815-0829 Osamu Segawa20009-0000-2469-6098 Takeshi Nakamura20009-0001-4991-3383 Kazuhiro Hotta10000-0002-5675-8713

    September 9, 2024
==============================================================================================================================================================================

§ ABSTRACT 

Named Entity Recognition (NER) encounters the challenge of unbalanced labels, where certain entity types are overrepresented while others are underrepresented in real-world datasets. This imbalance can lead to biased models that perform poorly on minority entity classes, impeding accurate and equitable entity recognition.


This paper explores the effects of unbalanced entity labels of the BERT-based pre-trained model. We analyze the different mechanisms of loss calculation and loss propagation for the task of token classification on randomized datasets. Then we propose ways to improve the token classification for the highly imbalanced task of clinical entity recognition.





§ INTRODUCTION


Named entity recognition (NER) is a Natural Language Processing (NLP) task of identifying and categorizing entities (such as names, events, things, and places) in a given raw text. An entity may span over just a single word or many continuous words. This process facilitates the automated analysis, search, and organization of extensive text datasets. One of the aims for NER is achieving higher precision for the entity labels with a relatively small number of training samples <cit.>. This creates the problem of imbalance between the entities that are recognized <cit.>. 

Within the clinical domain, NER holds particular significance, as it operates amidst a lexicon rich with specialized terminology necessitating accurate interpretation while imposing strict tolerances for errors. Capitalizing on the efficacy of transformer-based language models <cit.>, particularly BERT, has emerged as a dominant deep learning model for NER tasks. These models thrive on their foundational unsupervised pre-training on vast textual data, enabling them to encapsulate intricate linguistic structures that lend themselves to various language processing tasks.


The challenge of imbalanced labels considerably complicates the process of fine-tuning transformer models for Named Entity Recognition (NER). The distribution of entity classes often exhibits a significant skew, leading to disparities in the frequency of different entity types. This scenario frequently results in the underrepresentation of certain entity categories and the overrepresentation of others. Consequently, this inequality poses a hurdle to the model's capacity to generalize effectively to novel and unobserved data, particularly when it comes to the identification of essential yet infrequent entities, such as critical clinical information. In the realm of NER, it becomes imperative to address the issue of unbalanced labels, as it is pivotal to the development of models that ensure precise and equitable recognition of entities across all classes. This endeavor subsequently enhances the comprehensive efficacy and dependability of NER systems within a diverse array of practical applications.



This work makes a twofold contribution. Initially, we introduce a novel empirical bias testing methodology for BERT in token classification. We analyze the implications stemming from the application of arbitrary labels in BERT training. Furthermore, drawing on the observations and insights derived from related studies, we present a binary token labeling approach aimed at mitigating biases unsupported by empirical evidence. This enhancement seeks to augment BERT's capability to accurately discern entities characterized by a relatively low number of samples in contrast to the entities prevailing in the majority class.












 §.§ Related Works


The proliferation of text-based data in the biomedical field, such as electronic health records, clinical documents, and pharmaceutical specifications, has led to the widespread adoption of deep learning and Natural Language Processing (NLP) methods for extracting and processing information <cit.>. Additionally, studies have demonstrated that language models can partially encode clinical knowledge <cit.>. Contemporary generalized large-scale language models, which represent the forefront of language technology, exhibit suboptimal performance when deployed in clinical contexts, consequently undermining their reliability for clinical text analysis, as extensively noted in recent scholarly contributions <cit.>.
Therefore, biomedical and clinical NLP pose unique challenges, particularly the need to integrate structured domain knowledge into text representations, which is less prevalent in other domains <cit.>. To ensure the reliability of neural language modeling in the specialized medical field, models must learn directly from domain-specific terminologies rather than solely relying on general text data. As a result, significant research efforts within the medical NLP community have been devoted to integrating information from knowledge graphs into language models <cit.>. As a consequence, NER retains its position as the prevailing technique for clinical text analysis. However, it is noteworthy that the utilization of NER for clinical texts introduces a noteworthy challenge in the form of unbalanced accuracies. This imbalance in performance is intricately linked to the disparate distribution of data across distinct entity categories <cit.>.


The recent advancements in the field of biomedical NER, as highlighted in previous studies <cit.>, predominantly revolve around a restricted set of named entities such as diseases, chemicals, and genes. Nonetheless, it becomes imperative to broaden the scope of consideration to encompass a broader spectrum of biomedical entities. This includes entities pertinent to clinical diagnoses like diseases, symptoms, medical terms, risk factors, and vital signs, as well as epidemiological entities like infectious diseases and patient demographic information.

 

Certain studies in the biomedical field have explored the applicability of BERT in tasks related to biomedical NER <cit.>, yielding remarkable levels of performance.  An analysis of BioBERT, a transformer-based model specifically refined through fine-tuning procedures within the clinical text domain, has unveiled a pivotal insight. It discerns that the principal factor contributing to erroneous inferences generated by the BioBERT model resides in its limited grasp of the domain-specific knowledge <cit.>. 





§ BLACKBOX ANALYSIS FOR EMPIRICAL UNCERTAINTY PERSISTENCE



BERT (Bidirectional Encoder Representations from Transformers), introduced by , represents a profound architecture grounded in self-attention mechanisms. It undergoes a pretraining phase utilizing substantial volumes of data, guided by a language modeling objective. This model yields intricate linguistic text representations that have exhibited their utility across a multitude of tasks within the domain of natural language processing. Since its inception, BERT has undergone meticulous examination and practical application across diverse domains <cit.>. One of the downstream tasks for BERT is token classification, i.e., to use its contextualization capability to identify labels for words. BERT produces a 768-dimension vector for each token, processed to take into account a small amount of information about each of the other tokens in the input text. A downstream layer of a neural network can then learn to classify each token into entity categories. During fine-tuning, BERT's pre-trained weights and the final classification layer are modified to match the target task's label set. This process allows BERT to capture intricate contextual information from the input text and refine its predictions according to the token-level classification objectives, leading to robust and state-of-the-art performance across a diverse range of token-level classification tasks. In the case of NER, these labels identify types of entities that are learned from a particular annotated dataset. 

Fine-tuning BERT with unbalanced token labels is a significant consideration when adapting the model to token classification tasks where certain classes are disproportionately represented in the dataset. In scenarios where some classes occur much less frequently than others, the standard fine-tuning process can lead to biased models that perform well on majority classes but struggle with minority classes. The Unbalanced labels pose a significant challenge, as skewed distributions of entities in real-world datasets lead to accuracy issues, particularly for underrepresented classes. This imbalance affects model generalization and can hinder accurate recognition of minority entities, necessitating solutions for equitable and effective entity recognition in various applications.



 §.§ Maximum Likelihood Dilemma


Because NER datasets can have limited labelling and significant imbalances, we contend that the conventional cross-entropy loss function, while theoretically capable of asymptotically generating the optimal token classifier based on the maximum likelihood principle, would not succeed in delivering satisfactory performance in situations characterized by imbalanced data distributions. In simpler words, if there is a training set comprising 99% of class `O' and the remaining 1% of class `M', there is little incentive for the optimizer to optimize in favour of class `M' if it comes at the cost of a loss for class `O'. Half of this problem is solved by using weighted cross-entropy loss. However, when a calculated loss is back-propagated, the optimizer is likely to mostly adjust the weights that were modified by the majority class because the majority class had a bigger share in the gradient during the forward pass. The layers of the neural network are oblivious to loss at the individual token level because the loss is calculated for the whole batch that contains at least one sentence. 











 §.§ Exaggerated Empirical Bias in BERT for Token Classification


We test the above-mentioned argument on a clinical entity dataset MACCROBAT has has 41 annotated entities ranging in occurrences from 10 to 1208 in 200 clinical documents <cit.>. We use 85% of this annotated data to fine-tune BioBERT v1.1 model for 20 epochs for the token classification task as shown in Fig. <ref>. Then we use the rest of 15% of the documents to create a histogram of logits distributions and calculate the percentage of all the (TP+FP) positive predictions (represented by A) for each class as shown in Fig. <ref>.
















 §.§ Testing BERT with Arbitrary Token Labels
 

Assuming that there is inherent arbitrariness to language <cit.>, certain factors in languages can't be learned as hard and fast rules. We argue that if there is fuzziness in the training corpus, a good language model should retain the fuzziness if there is no significant evidence for clarity. We test this assumption on BERT by fine-tuning it on randomly generated token labels for the same clinical text corpus. We replace the original labels with randomly generated 3 classes of labels, 60% of class `O', and 20% each for classes `M' and `N' as shown in Fig. <ref>. Since the labels are randomly assigned, the model should not learn any significant pattern other than the unbalanced amount of labels. We train the BERT-base-cased model for 30 epochs using 85% of the corpus and measure the logits distributions on the rest 15% of the corpus after each epoch. We test the evidence-based uncertainty of the fine-tuned model by comparing the number of labels for each class N with the number of predicted positive labels for each class represented as A. The difference between the two numbers will be a measure of the empirical dependency of the model since there is no other pattern to learn from other than the empirical amount of labels. Figure <ref> shows the logits distributions of 3 classes for the test set. It can be observed that the model has an unaccounted bias towards the class `O' because it predicts the highest logit value for class `O' 87.8% of the time even though empirically it should be 59.9%. 









§ BINARY LABELS FOR TOKEN CLASSIFICATION


When the loss is calculated by all token labels, it gets diluted into a generic loss that does not vary much from batch to batch. The problem of deeper layers being oblivious to the non-variant loss at the final layer can be solved by having only two classes creating a higher variation in loss so that it can backpropagated deeper into the model <cit.>.  We theorize that if a contextualized token classification model is fine-tuned to recognize only two labels (binary classification), then it is likely to learn each entity more precisely as compared to a model that is fine-tuned to recognize more than two labels within the same batch. We propose to use Binary Token Labels (BTL) to finetune BERT for the NER task as opposed to the conventional method of using all token labels (ATL) for each batch. The BTL method splits training batches into multiple copies but each batch of passages has binary contrastive labelling against `O' (the absence of any entity label is represented by `O') for only one true-positive label while the rest of the labels are masked. 







The weighted cross-entropy loss at the token level is calculated as

    l_n= -w_y_nexp(x_n, y_n)/∑_c=1^Cexp(x_n, c)

where x_n, c is the value of the logit at the output layer for class c, x_n, y_n is the logit value for the target class, C is the total number classes, w_y_n is the weight of the target class that is calculated as

    w_c = 1 - N_c/N

where N_c is the number of tokens for the class c and N is the number of all labelled tokens in the training set.

The mean loss L for the whole batch is

    L = ∑_n=1^Bl_n/∑_B^n=1w_y_n·{v_n ≠ X }



where B is the number of total tokens in the batch, and X is the masked label conditioned as

v_n={[             X y_n ∉{O, C_b};             1    otherwise. ].

The C_b is the label for batch b that is left unmasked along with the true negative label O. The difference from the usual approach for loss calculation is that each batch contains a mixture of all target classes. Whereas in this approach, each batch only has a true positive label for only one entity class. 




We performed the uncertainty persistence test on BERT-base-case again using the BTL approach for the randomly labelled dataset. It can be seen in Figures <ref> and <ref> that using BTL the model is less sensitive to the empirical bias as it converges to the maximum-likelihood over the epochs. This makes it possible to intervene in the training process before the maximum likelihood end goal of the optimizer is reached. This approach focuses on the core task of entity presence or absence, allowing the model to learn to distinguish entities from non-entities effectively. Lastly, binary fine-tuning can yield models that are less sensitive to label noise or annotation inconsistencies for a batch.
















 §.§ Experimentations


Based on the observations made on the randomly labelled dataset in Section <ref>, we tested the BTL approach on the MACCROBAT dataset with the intention that the loss caused by the majority classes won't undermine the loss caused by the small classes as it did in Figure <ref> when we tested the conventional token classification learning approach. The bigger documents in the MACCROBAT dataset are broken into smaller chunks to fit the maximum input tokens size of 512. This results in 200 documents being divided into 886 passages for the training set and 169 passages for the test set. Another few of the quantitative entity labels such as Volume, Mass, Height, and Weight are merged into a single category. 

To achieve a balanced performance across all entity classes we create a clinical NER method with the following measures:


    
  * Weighted cross-entropy using the class weights as calculated in Section <ref>.
    
  * The binary token label (BTL) approach is used to create segregated batches for each entity class. We also test the conventional all-token-labels (ATL) approach where each batch has all entity labels without any masking. 
    
  * The number of batches for each entity class is balanced. The minority class batches are repeated more frequently to match the number of batches of the largest class `O'.
    
  * Both models (BTL and ATL) are fine-tuned using pre-trained BioBERTv1.1. The training is run for 20 epochs using an SGD optimizer with a learning rate of 5e-5.
    
  * A latent KNN classifier (17 neighbours) to predict the final entity label using the raw logits from the output layer. The KNN classifier is trained separately after finetuning.




The utilization of K-Nearest Neighbors (KNN) serves the purpose of additional independent calibration post finetuning. The outcomes for 33 entities within the MACCROBAT dataset are presented in Table <ref>. Notably, employing the BTL method leads to a substantial increase in unweighted measures like unweighted accuracy and mean precision across all entities. This improvement is particularly pronounced for entities with small sample sizes.



Figure <ref> shows the distribution of logits for the test set using the BTL approach. The number of positive predictions for small-sample entity classes are much closer to empirically expected positive predictions as compared to the distributions for ATL model (Figure <ref>). Another observable difference from Figure <ref> is the wider spread of the distributions. This is due to the increase in the entity-specific loss variation as discussed at the start of this section.




In Table <ref>, a comparison of F1 scores between the proposed method and the baseline is provided. However, it's important to clarify that the baseline's F1-score represents the overall weighted F1 score. While examining performance metrics for select individual entities, a marked distinction from the baseline results by   becomes evident. The advantage of BTL is the increase in the significant increase in unweighted accuracy which is a better measure of balanced accuracy across all entities, however, it is rarely reported by other works.










§ LIMITATIONS

This study conducted an analysis of the empirical biases inherent in BERT's NER token classification. A novel black box testing method was introduced to assess empirical biases, independent of linguistic content, with the caveat that matching expectations through empirical evidence is valuable in domains with established parameters. Subsequently, insights gleaned from the black box testing were leveraged to enhance NER performance on a heavily imbalanced clinical dataset. While there was an overall enhancement in recognition metrics, it's worth noting that the performance improvements were not uniform across all entities. Certain entity labels, such as 'Coreference,' experienced a decrease in accuracy when utilizing the proposed approach.




§ ETHICAL STATEMENT

Our conducted experiments and the model framework we propose are designed to promote investigation within the clinical information extraction domain while prioritizing the prevention of privacy breaches. The data utilized in our study is publicly accessible and has been thoroughly de-identified. Although recent studies have demonstrated the challenge of reconstructing sensitive personal information from such data, a minimal potential risk exists for future models to achieve this. It's important to note that we have not made any modifications to the data's content that would enhance the probability of such an eventuality, thus ensuring the mitigation of any risks related to the leakage of private information.




acl_natbib




