WeLT: Improving Biomedical Fine-tuned Pre-trained Language Models with Cost-sensitive Learning
Year: 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 costsensitive fine-tuning approach based on new rescaled 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 reweighting 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.

text: The 22nd Workshop on Biomedical Natural Language Processing and BioNLP Shared Tasks, pages 427–438 July 13, 2023 ©2023 Association for Computational Linguistics

heading: 1 Introduction
text: 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 (Devlin et al., 2019; Lee et al., 2019). 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 (ValizadehAslani et al., 2022). Most of the real-world biomedical datasets
are highly imbalanced (Akkasi et al., 2018). 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 (Shen and Sarkar, 2005). 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. Table 1 shows the imbalanced nature of BioNER corpora for multiple entity types, including chemical, disease, gene, and species entities. It is clear to conclude that biomedical ground-truth training datasets are highly skewed. Such high data imbalance and scarcity make many BioNER models biased towards the O class, thus, they often misclassify entities (B and I classes).
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 (Lee et al., 2021). On the other hand, fine-tuning BioPLMs (Alrowili and Shanker, 2021) naively on highly skewed datasets
427
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:
1. We propose WeLT, a class-balanced reweighting scheme added to a trainer’s loss function before fine-tuning models.
2. We compare WeLT to vanilla fine-tuning approach, and existing cost-sensitive class weighting methods, including Inverse Number of Samples, Inverse of Square Root of Number of Samples, and Effective Number of Samples (Suri, 2022; Cui et al., 2019). We conduct our experiments on several transformers such as BERT and ELECTRA.
3. We publicly release the code and hyperparameters to reproduce our research results1. We also release all the fine-tuned models on the Hugging Face Hub2 (Wolf et al., 2020).

heading: 2 Related Work
text: Large pre-trained language models have been the state-of-the-art across a variety of tasks in natural language processing (Devlin et al., 2019; Clark et al., 2020). However, direct application of these models to biomedical text mining applications
1https://github.com/mobashgr/WeLT 2https://huggingface.co/mobashgr
yields unsatisfactory results, since the word distributions of general domain and biomedical corpora are very diverse (Lee et al., 2020). 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). BlueBERT (Peng et al., 2019) is pre-trained on PubMed abstracts and clinical notes. SciBERT (Beltagy et al., 2019) 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 (Gu et al., 2021). Their results show that domain-specific pretraining from scratch can substantially outperform the traditional mixed-domain approach. BioELECTRA is pre-trained on PubMed and PubMed Central (Kanakarajan et al., 2021). Biomedical BERT variants and BioELECTRA require minimal modification of the architecture. Leaman and colleagues developed various ensemble methods for the chemical identification task (Leaman et al., 2023) 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 (Devlin et al., 2019; Lee et al., 2019; Dodge et al., 2020), and (c) optimization difficulties early in training (vanishing gradients) (Mosbach et al., 2020). 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 (Zhao et al., 2018).
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 (Akkasi et al., 2018). 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 (Elkan, 2001). 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 (Suri, 2022). Cui and colleagues proposed a reweighting loss by using the inverse effective number of samples for addressing class imbalance (Cui et al., 2019).
Handling class imbalance for different NLP downstream tasks is challenging, due to the following reasons:
• Most of the training data are part of gold standard datasets that are annotated and curated by domain experts. The human annotation and curation of biomedical data are tedious tasks that require manpower and time, therefore, there are only limited biomedical goldstandard datasets. Consequently, undersampling is not the appropriate solution to address class imbalance.
• Even if the training data is not a gold standard, transformers are powerful enough to memorize duplicate instances. Moreover, recent research proves that de-duplicating training data makes the language models better and allows faster training (Lee et al., 2021).
• Removing some instances from the majority classes can result in information loss.
• Running multiple experiments to acquire new class weights by applying cost-learning approaches is a tedious task that does not provide a generic method that can be used to acquire new weights for both majority and minority classes.
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.

heading: 3 WeLT - Weighted Loss Trainer
text: 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).
wx = 1− Ix It , with x ∈ {o, b, i} (1)
We illustrate the computation based on the NCBI dataset. Let TDncbi be the training data of NCBI, consisting of It = 5,868,591 instances of the three classes o, b, and i. The number of instances per class in TDncbi is denoted Io, Ib and Ii for the O, B, and I classes, respectively.
The three class instances are as follows:
• Io = 4,262,718
• Ib = 389,892
• Ii = 1,215,981
The new weight for each of the three classes is computed using Equation (1), leading to the following values :
• wo ≈ 0.27
• wb ≈ 0.79
• wi ≈ 0.93
Let W⃗V be the vector that has three elements, i.e.
WV =     wo wb wi    . Then, W⃗V is normalized using Softmax, as shown in Equation (2).
σ ⃗(WV )i = eWVi∑t c=1 e WVc (2)
where W⃗V is the input vector to Softmax. WVi are the elements of W⃗V , so WVi of TDncbi contains three elements.
The standard exponential function, denoted eWVi , is applied to each element of the input vector W⃗V .
The normalization term, denoted ∑t
c=1 e WVc ,
ensures that all the output values of the function will sum to 1, where t = 3.
Thus, the values for normalized vector of
TDncbi denoted, σ ⃗(WV ), are ≈
   
0.2167 0.4192 0.3641
   
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.
Lce = − c∑
i=1
yi log ŷi (3)
where c = 3 represents the different classes to be predicted, yi ∈ [0, 1] denotes the ground truth class, and ŷi ∈ [0, 1] is the model’s estimated probability for each class with ground truth. Since the biomedical gold standard training datasets are highly imbalanced, we use the weighted cross-entropy loss.
The normalized vector σ ⃗(WV )i is passed to the weighted cross-entropy loss function, as illustrated in Equation (4), after extending the class trainer and overriding the compute_loss function (Paszke et al., 2017).
Lwce = σ ⃗(WV )i Lce (4)
Subsequently, the models are fine-tuned using WeLT with the exact training cost as the vanilla fine-tuned approach.

heading: 4 Experiments
text: 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.

heading: 4.1 Evaluation Datasets
text: NCBI Disease The American National Institutes of Health released the NCBI disease corpus to promote disease NER research. The public release of the NCBI disease corpus contains 6,892 disease mentions, which are mapped to 790 unique disease concepts.
BC5CDR-Disease and BC5CDR-Chemical The BioCreative V Chemical Disease Relation (CDR) corpus was created for the Chemical Disease Relation (CDR) Task. It consists of human annotations of all chemicals, diseases, and their interactions in 1,500 PubMed articles.
BC4CHEMD The BioCreative IV Chemical and Drug (BC4CHEMD) named entity recognition task corpus. It contains 10,000 abstracts annotated for mentions of chemical and drug names.
BC2GM The BioCreative II Gene Mention task corpus. It consists of 20,000 sentences from biomedical publication abstracts, annotated genes, and proteins.
Linnaeus Linnaeus corpus has 100 full-text documents for species annotations.
BioRED-Disease and BioRED-Chemical The BioRED corpus was created for multiple biomedical relations. It consists of human annotations of all different biomedical entities and their interactions in 600 PubMed abstracts. In this work, we only focus on chemical and disease instances.

heading: 4.2 Evaluation Metrics
text: Regarding the evaluation metrics, we report the entity-level precision, recall, and F1 scores. However, we assess the performance based on the entitylevel F1 score. Due to the nature of our work that investigates the impact of addressing the class imbalance before fine-tuning, we do not compete with the state-of-the-art BioNER baselines. However, we compare the vanilla fine-tuning approach to WeLT and three existing weighting schemes using the same hyper-parameters.
Besides the overall assessment of the aforementioned fine-tuned models, we evaluate the annotation quality for species entities on the Linnaeus dataset. Thus, we used two sequence labeling metrics: (a) seqeval (Nakayama, 2018), and (b) FairEval (Ortmann, 2022). FairEval is one of the latest metrics on which Ortmann argues that the traditional evaluation metric causes double penalties for close-to-correct annotations. Therefore, Ortmann developed FairEval, which ensures that every
error is counted only once.

heading: 4.3 Implementation
text: We adapt the BioBERT (Lee et al., 2020) PyTorch named entity recognition code to develop WeLT. In our experiments, we used the five following pre-trained model variants (a) BioBERT, (b) BlueBERT, (c) PubMedBERT, (d) SciBERT, and (e) BioELECTRA. For more details on the hyperparameters, see Appendix A. All the experiments were carried out using two Tesla P40 GPUs with 24GB memory.

heading: 4.4 Compared Methods
text: As previously mentioned, we avoid using traditional resampling approaches, as there might be a possibility of information loss, and duplicating training examples leads to poor performance of language models (Lee et al., 2021). Therefore, in this paper, we focus on existing weighting schemes.
We compare WeLT to a vanilla fine-tuned approach (i.e., without handling class imbalance). We also evaluate WeLT with other weighting schemes: (a) Inverse of Number of Samples (INS), (b) Inverse of Square Root of Number of Samples (ISNS) , and (c) Effective Number of Samples (ENS) (Suri, 2022; Cui et al., 2019). Regarding the ENS approach, we used different values for β as follows: (a) 0.3, (b) 0.5, and (c) 0.9. For more details about the aforementioned weighting schemes, see Appendix B.

heading: 5 Results and Discussion
text: 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 BioPLMs 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,
(b) pre-training approach, and (c) class imbalance percentage.
According to the statistics in Table (11), BioRED is the smallest dataset. WeLT has the best performance for fine-tuning experiments except for PubMedBERT on BioRED-Disease, as shown in Table (3). The same applies to BioRED-Chemical in Table (4) except for BioBERT and BlueBERT. On the other hand, BC4CHEMD is the largest dataset. WeLT achieves the second-best score for all the experiments in Table (7), except for BioBERT.
Regarding the class distribution percentage in Table (1), Linnaeus is the highest skewed dataset, and NCBI is the least imbalanced one. WeLT has the highest score for the experiments related to Linnaeus as shown in Table (2), except for the secondbest score while fine-tuning BioBERT. Similarly, WeLT has the best score for all the experiments using NCBI except for BlueBERT, as shown in Table (9).
WeLT has the best score for the BC5CDRChemical experiments in Table (5) except for finetuning SciBERT and the second-best score for PubMedBERT. BC5CDR-Disease experiments in Table (6) show that WeLT has the best score for all the experiments, except fine-tuning SciBERT, and got the second-best score for BioBERT. Finally, BC2GM experiments in Table (8) show that WeLT outperformed except for being second-best score for fine-tuning SciBERT.
We believe that WeLT shows advantages over the chosen existing weighting schemes, because it considers the overall dataset class distribution. None of these weighting schemes considers the overall
class distribution while calculating the re-scaled weights. WeLT presents a generic approach to calculate re-scaled weights based on their normalized inverse effect in the dataset. In other words, major and minor classes have smaller and higher weights, respectively, in WeLT. This gives higher loss costs for minor classes while fine-tuning.
We highlight special patterns and give some intuitions about the successful and failed cases related to the performance of WeLT:
1. Fine-tuning WeLT on the largest dataset (BC4CHEMD) has shown the second-best score for all the experiments but failed for finetuning BioBERT. Based on our observation, ENS variants have the best performance. We believe that considering the overall class distribution for calculating the re-scaled weights degrades the performance. Further investigations should be considered by adding data size as an additional factor. Despite ENS variants outperformance, we emphasize that adding extra hyper-parameters in ENS is problematic and expensive since the appropriate β factor is unknown.
2. WeLT has the best fine-tuning scores for the highly skewed dataset (Linnaeus) on all the experiments, except for BioBERT. We believe that the calculation of new re-scaled weights has a positive impact on performance. The rest of the other weighting schemes focus only the number of class samples and not on the overall class distribution.
3. Fine-tuning BioELECTRA using WeLT has
the best scores except for BC4CHEMD. BioELECTRA is the biomedical version of ELECTRA. ELECTRA’s pre-training strategy uses a more efficient approach called "replaced token detection". It learns from all input tokens rather than just the small subset that was masked out like in BERT models. It uses another neural network that aims to trick the model by replacing random tokens with fake tokens.
4. Regarding fine-tuning PubMedBERT as domain-specific pre-trained language model, WeLT has the best score for these four datasets as follows: (a) Linnaeus, (b) BioREDChemical, (c) BC5CDR-Disease, and (d) BC2GM. ENS variants has the best score for fine-tuning BioRED-Disease dataset. For the three other remaining datasets, WeLT has the second-best score. We believe that WeLT and different variants of ENS should be considered when fine-tuning PubMedBERT. The same applies for SciBERT experiments.

heading: 6 Error Analysis
text: 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:
• Type-1: an entity predicted by the NER model but not annotated in the gold-standard datasets. For instance, “S“ is detected by BioPLMs since it is an abbreviation for “Sulphur“. However, it was not annotated by human experts in the BC5CDR gold standard dataset.
• Type-2: an entity annotated in the goldstandard datasets but not predicted by the
NER model. The main issue behind the misclassification are abbreviated entities. For example, “PAN“, which is an abbreviation for “Peroxyacetyl nitrate“, is not recognized as a chemical entity.
• Type-3: an entity correctly predicted but has overlapping spans errors. For example, BioPLMs recognizes two chemical entities separately such as “amphotericin B-“ and “sodium deoxycholate“. However, the gold-standard annotation is “amphotericin Bsodium deoxycholate“.
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.

heading: 7 Conclusion and Future Directions
text: 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 BioPLMs, 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.
Limitations
We propose WeLT as an approach to address the class imbalance problem. We have only evaluated WeLT for BioNER tasks, although we hypothesize that it can be adapted to any application / domain that has skewed dataset. We further want to point
out that our method was only evaluated on English datasets. Yet, we argue that it can be applied to other languages as well. Finally, we have not assessed WELT’s performance on larger training datasets (i.e., more than 30,683 training examples).
Ethics Statement
We see no impending risk in the development of our proposed and potential applicability to other nonbiomedical applications. However, we would like to point out the aforementioned limitations (see the previous section). We report all the required hyperparameters, implementation details, and hardware specifications to reproduce our results. We release all the fine-tuned models on the Hugging Face Hub.

heading: Acknowledgments
text: WM and GM wish to acknowledge base funding by the HITS and the Klaus Tschira Foundation, KTS, supporting this work. This paper is supported by European Union’s Horizon 2020 research and innovation program under the Marie SkłodowskaCurie grant agreement PoLiMeR, No 812616

heading: A Hyper-parameters
text: We report all the hyper-parameters that were used to train the 280 models in the WeLT GitHub Repository.

heading: B Comparable Weighting Schemes
text: Inverse of Number of Samples (INS)
weight[class] = 1
nc (5)
where nc is the corresponding number of samples in the class.
Inverse of Square Root of Number of Samples (ISNS)
weight[class] = 1√ nc
(6)
The main difference between INS and ISNS is that INS increases the recall by decreasing the values of false negatives while having low precision as the weights of majority class has been diminished. Therefore, ISNS overcomes this by taking the square root of class frequencies so that the class weights would be larger than INS.
Effective Number of Samples (ENS)
weight[class] = 1− β 1− βnc (7)
where β is a hyper-parameter ∈ [0, 1].

heading: C Dataset statistics and other results
text: Due to page limitations, we report the rest of our experimental results and the gold-standard datasets statistics here.

heading: Future Work
text: 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. Limitations
We propose WeLT as an approach to address the class imbalance problem. We have only evaluated WeLT for BioNER tasks, although we hypothesize that it can be adapted to any application / domain that has skewed dataset. We further want to point
out that our method was only evaluated on English datasets. Yet, we argue that it can be applied to other languages as well. Finally, we have not assessed WELT’s performance on larger training datasets (i. e. , more than 30,683 training examples)