W e LT : Improving Biomedical Fine-tuned Pre-trained Language Models with Cost-sensitive Learning (2023)

Abstract: Fine-tuning biomedical pre-trained language models (BioPLMs) such as BioBERT has become a common practice dominating leaderboards across various natural language processing tasks. Despite their success and wide adoption, prevailing fine-tuning approaches for named entity recognition (NER) naively train BioPLMs on targeted datasets without considering class distributions. This is problematic especially when dealing with imbalanced biomedical gold-standard datasets for NER in which most biomedical entities are underrepresented. In this paper, we address the class imbalance problem and propose WeLT, a cost-sensitive fine-tuning approach based on new re-scaled class weights for the task of biomedical NER. We evaluate WeLT’s fine-tuning performance on mixed-domain and domain-specific BioPLMs using eight biomedical gold-standard datasets. We compare our approach against vanilla fine-tuning and three other existing re-weighting schemes. Our results show the positive impact of handling the class imbalance problem. WeLT outperforms all the vanilla fine-tuned models. Furthermore, our method demonstrates advantages over other existing weighting schemes in most experiments.

## introduction fine-tuning biomedical pre-trained language models has become a mainstream practice dominating leaderboards across various biomedical nlp tasks. despite major advancements and wide usage, recent work reported fine-tuning instabilities for general-domain nlp tasks @xcite . lee and colleagues reported that small training datasets (i.e., less than 10,000 examples) is one of the potential reasons for finetuning instabilities. moreover, a few studies point out the positive impact of handling the class imbalance before fine-tuning @xcite . most of the real-world biomedical datasets are highly imbalanced @xcite . nevertheless the impact of class imbalance before finetuning biomedical datasets is often explored, especially not for named entity recognition. the commonly used tagging scheme for biomedical named entity recognition (bioner) is the inside-outside-beginning (iob) format @xcite . it consists of three classes: (a) the b tag represents the beginning or first token of a biomedical entity, (b) the i tag denotes the continuation of the first token as an inside biomedical entity, and (c) the o tag represents a token that is not part of a biomedical entity. thus b and i classes are the positive samples, while the o class is the negative sample. since these gold-standard biomedical datasets are curated by domain experts, we avoided using traditional resampling approaches, as there might be a possibility of information loss, and duplicating training examples leads to poor performance of language models @xcite . on the other hand, fine-tuning bioplms @xcite naively on highly skewed datasets is problematic, since it can lead to a bias in the trained models, which can negatively affect the performance. therefore, we investigate the impact of handling the class imbalances while fine-tuning. we develop a weighted loss trainer (welt) that addresses class imbalances by introducing coefficients that penalize majority classes and give more weight to the rare ones before fine-tuning. to evaluate our approach, we conduct extensive bioner experiments on eight biomedical goldstandard datasets. we fine-tune different variants of biomedical bert models (mixed/continual training and domain-specific) and bioelectra. we do not only compare welt against vanilla finetuning approach (i.e., it assumes that the misclassification errors are equal for all class labels), but also with three other existing comparable weighting schemes. welt and other weighting schemes outperform vanilla fine-tuned models. this proves the hypothesis that class imbalance has a negative impact on fine-tuning. we suggest that class imbalance should be considered as an additional factor for fine-tuning. in summary, our main contributions are as follows: we also release all the fine-tuned models on the hugging face hub foot_1 @xcite . ## related work large pre-trained language models have been the state-of-the-art across a variety of tasks in natural language processing @xcite . however, direct application of these models to biomedical text mining applications yields unsatisfactory results, since the word distributions of general domain and biomedical corpora are very diverse @xcite . therefore, lee and colleagues adapted bert for biomedical applications and pre-trained bert on pubmed abstracts and pubmed central full-text articles. however, biobert is a mixed-domain pre-training approach, as authors used the original vocabulary of bert (i.e., wikipedia and bookcorpus). blue-bert @xcite is pre-trained on pubmed abstracts and clinical notes. scibert @xcite generates the vocabulary and pretrains from scratch, using biomedical and computer science literature. subsequent work has been devoted to train the model from scratch on a biomedical vocabulary only, such as pubmedbert @xcite . their results show that domain-specific pretraining from scratch can substantially outperform the traditional mixed-domain approach. bio-electra is pre-trained on pubmed and pubmed central @xcite . biomedical bert variants and bioelectra require minimal modification of the architecture. leaman and colleagues developed various ensemble methods for the chemical identification task @xcite mostly focusing on fine-tuning biomedical bert variants. the fine-tuning instability of bert has been pointed out in various studies. the observed instabilities are identified due to the following potential reasons: (a) catastrophic forgetting, (b) small size of the fine-tuning datasets @xcite @xcite , and (c) optimization difficulties early in training (vanishing gradients) @xcite . one of the possible reasons for poor fine-tuning results is using an imbalanced dataset even if the training dataset is large. most of the biomedical datasets are skewed @xcite . most existing algorithms for learning imbalanced datasets can be divided in to two categories: (a) re-sampling, and (b) re-weighting. in the biomedical domain, akkasi and colleagues directly investigated the class imbalance problem for biomedical datasets. they propose a balanced undersampling approach for sequence data and enhance the classification performance by systematically removing negative samples from training data @xcite . re-weighting approaches allow minority samples to contribute more to the loss function. one of the earliest approaches introduced a factor that can be multiplied by a certain thresh-old, resulting in higher weights for misclassification of minority classes @xcite . however, this factor assumes that one deals with a binary classification problem. recent work addresses class imbalance when applying bert for sentence classification. for example, madabushi and colleagues (madabushi et al., 2020) applied cost-weighting for a binary classification problem on which the exact weight is related to the dissimilarity of training, development, and test datasets. they report that the weights are obtained experimentally through a hyper-parameter search. weighting by inverse class frequency or a smoothed version of inverse square root of class frequency are often adopted @xcite . cui and colleagues proposed a reweighting loss by using the inverse effective number of samples for addressing class imbalance @xcite . handling class imbalance for different nlp downstream tasks is challenging, due to the following reasons: based on the extensive literature survey related to the existing class imbalance research, we conclude that the cost-sensitive learning approaches for biomedical bert-based models have been neglected for bioner. ## welt -weighted loss trainer vanilla fine-tuned models assume that misclassification errors and cost, respectively, contribute equally to all class labels. however, in many realworld applications, such as bioner, the differences between different misclassification errors can be quite large and thus critical. in other words, the error costs for rare classes in a trainer's loss function should be higher. this can be achieved by re-weighting the instances of each class according to their misclassification costs. for this, we proposed welt, a cost-sensitive fine-tuning approach for bioner as depicted in figure 1 . welt sets a higher class weight for the minority classes and a lower weight for the majority classes. based on the context of our work with the iob tagging scheme, o is the majority class. while b and i are the minority classes as shown in figure 2 . we present a generic way of calculating the re-scaled class weights for the iob tags without having hyper-parameter search factors. the new class weights are calculated using the inverse ratio of each corresponding class distribution over the total class distributions of the training datasets, as shown in equation ( 1 ). we illustrate the computation based on the ncbi dataset. let t d ncbi be the training data of ncbi, consisting of i @xmath0,868,591 instances of the three classes o, b, and i. the number of instances per class in t d ncbi is denoted i o , i b and i i for the o, b, and i classes, respectively. the three class instances are as follows: the new weight for each of the three classes is computed using equation (1), leading to the following values : let ⃗ w v be the vector that has three elements, i.e. the standard exponential function, denoted e w v i , is applied to each element of the input vector ⃗ w v . the normalization term, denoted t @xmath1, ensures that all the output values of the function will sum to 1, where @xmath2. as previously mentioned vanilla fine-tuned trainers use cross-entropy loss. cross-entropy loss as shown in equation ( 3 ) minimizes the training error by assuming that individual samples and classes are equally important as if the class frequencies are sufficiently balanced. subsequently, the models are fine-tuned using welt with the exact training cost as the vanilla fine-tuned approach. ## experiments we empirically demonstrate the performance of our method through various experiments on eight biomedical gold-standard datasets focusing on the bioner task. we evaluate welt on both mixeddomain and domain-specific pre-trained models. we compared welt to their corresponding vanilla fine-tuning approach and three existing weighting schemes. we assess the behaviour of welt when being fine-tuned while dealing with different dataset sizes and a variety of class distributions. in addition to the experimental analysis, we further share the implementation details and evaluation metrics. ## results and discussion we present the results of seven fine-tuning experiments that include (a) ins, (b) isns, (c) ens with β = 0.3, (d) ens with β = 0.5, (e) ens with β = 0.9, (f) vanilla, and (g) welt. for a fair comparison of all eight biomedical datasets, we applied the same experimental settings for all the fine-tuning approaches. we fine-tuned five bio-plms using six weighting schemes and a vanilla approach on eight targeted datasets. thus, we have 280 experimental results as presented in tables (2)-( 9 ). at first glance, addressing the class imbalance enhances the performance. welt, different variants of ens, and isns have the highest f1-score compared to their corresponding vanilla fine-tuning approach. we extensively report the results based on the three following criteria: (a) size of training datasets, we highlight special patterns and give some intuitions about the successful and failed cases related to the performance of welt: ## error analysis we have further analyzed different types of bioner mismatches for error analysis. based on our observations, we noticed that still trivial bioner errors occur. yet, addressing the class imbalance before fine-tuning has a positive impact on the overall performance and the sequence labeling evaluation as shown in table ( 10 ). as a proof-ofconcept, we further evaluated the tagging quality outputs of each fine-tuning approach on the linnaeus datasets. we report the f1-scores using seqeval with strict mode and faireval with fair mode. welt has the best score for fine-tuning bluebert, scibert and bioelectra fine-tuned models and got the second best score for biobert and pubmedbert. the observed mismatches occur due to the following three types of errors: we believe that the first two mismatch error types require knowledge enrichment and usage of active learning approaches to tackle such semantic annotation issues. the third mismatch error can be fixed by enhancing the post-processing method to better merge tokens that are part of a recognized entity. ## conclusion and future directions in this work, we propose a weighted loss trainer (welt) as a cost-sensitive approach that addresses the class imbalance problem of biomedical goldstandard datasets before fine-tuning models. we have evaluated welt using five different bio-plms, including general-domain and domainspecific pre-training. we do not only focus on biomedical bert variants, but also bioelectra. we have conducted 280 experiments in total (8 datasets, 5 bioplms, and 7 different fine-tuning approaches). furthermore, we have extensively evaluated welt's performance against other finetuning approaches. moreover, we present a thorough error analysis. we only focus on bioner as a downstream task but we believe that welt can be adapted to other tasks with a class imbalance problem. welt can also be applied to non-biomedical domain applications. in the future, we plan to evaluate welt's performance for fine-tuning larger datasets (i.e., more than 30,683 training examples). more precisely, in this work, bc4chemd is the largest dataset with 30,683 training instances.