Assessment of D istil BERT performance on Named Entity Recognition task for the detection of Protected Health Information and medical concepts (2020)

Abstract: Bidirectional Encoder Representations from Transformers (BERT) models achieve state-of-the-art performance on a number of Natural Language Processing tasks. However, their model size on disk often exceeds 1 GB and the process of fine-tuning them and using them to run inference consumes significant hardware resources and runtime. This makes them hard to deploy to production environments. This paper fine-tunes DistilBERT, a lightweight deep learning model, on medical text for the named entity recognition task of Protected Health Information (PHI) and medical concepts. This work provides a full assessment of the performance of DistilBERT in comparison with BERT models that were pre-trained on medical text. For Named Entity Recognition task of PHI, DistilBERT achieved almost the same results as medical versions of BERT in terms of F1 score at almost half the runtime and consuming approximately half the disk space. On the other hand, for the detection of medical concepts, DistilBERT’s F1 score was lower by 4 points on average than medical BERT variants.

## introduction clinical records play an important role in the discovery of disease treatment and the advancement of medical research @xcite . the clinical text corpora used for research includes doctor's notes, clinical study reports and medical articles. there are several regulations that control the use and transfer of personal information such as general data protection regulation (gdpr) in europe, personal information protection and electronic documents act (pipeda) in canada and health insurance portability and accountability act (hipaa) in the us. hipaa safe harbor for example lists 18 attributes that can potentially identify an individual and dictates that all of them need to be de-identified before a dataset can be shared for secondary use such as research @xcite . one possible approach is the manual annotation and de-identification of clinical text. this approach is simply not feasible due to the high cost of experts manually annotating clinical documents @xcite . due to the advancement of natural language processing research, the deidentification of phi was framed as a named entity recognition (ner) problem that can be solved by deep learning techniques. this work fine-tuned a deep learning model on a medical corpus and assessed its quality in detecting phi and medical concepts in comparison with models whose embeddings were generated from a medical corpus. the paper is organized as follows: in the next section we review the state-of-the-art for solving ner tasks used in the detection of phi. in section 3 and 4 we define the problem and detail our methodology. in section 5 we present our results and we finally conclude in section 6. ## related work in a major breakthrough in nlp research, a simpler neural network architecture was introduced by @xcite called transformers which is an attention-based mechanism. its main premise was to do away with recurrence and convolution in neural networks. self attention generates a representation by connecting different positions of a given sequence. self attention is easier to parallelize and enables better understanding of longrange dependencies. transformers enabled the introduction of bidirectional encoder representations from transformers (bert) by @xcite . bert allows the generation of representations utilizing context from both directions of a sequence. it consists of two steps: pre-training and fine-tuning. pre-training is the unsupervised learn-ing step to generate the representations. bert was pre-trained on wikipedia and bookcorpus. the pre-training step consists of two tasks: masked language model (mlm) which masks a certain percentage of the input sequence and attempts to predict those missing tokens. the second pretraining task is next sentence prediction (nsp) which was specifically added to help with tasks involving relationship between a pair of sentences such as question answering. the fine-tuning is the supervised learning portion where bert is trained on custom datasets by the user for their respective tasks with little to no feature engineering required for a specific nlp downstream task. further attempts have been made to improve on the original bert such as roberta introduced by @xcite which assessed the impact of different hyperparameters and concluded that training over longer sequences and removing nsp achieves better results. other bert variations were also pre-trained on medical domain corpora such as biobert @xcite , bluebert @xcite and clinicalbert @xcite that were pre-trained on pubmed which contains biomedical research articles and mimic-iii which contains doctors' notes from the intensive care unit admissions. the medical versions of bert proposed higher f 1 score performance when evaluated on biomedical tasks including ner. bert and its variations, however, require extensive computational resources to deploy in production environments. to address these limitations, distilbert was introduced by @xcite . the authors applied the concept of knowledge distillation to produce a lighter version of bert that is 40% smaller, 60% faster and achieves 97% of the original bert f 1 score when measured on question answering task. it can also be deployed on lower power computing chips such as mobile devices to run predictions. further studies were published to assess the performance of distilbert compared to other state-of-the-art models. in a study by @xcite , it compared dis-tilbert's performance against elmo on two text classification tasks. the first was a binary classification task of protest and non-protest news from english articles from local newspapers in india and china. the second task was a sentence classification task of movie reviews on rotten tomatoes. the authors concluded that distilbert generalizes better than elmo while having similar f 1 score. wang et al. (2020) also used distilbert for a machine translation task to generate synthetic data to diagnose language impairment in children. dis-tilbert achieved 5% and 15% higher f 1 scores when compared with elmo and word2vec respectively. ## problem statement although bert achieved state-of-the-art for a wide variety of nlp tasks, they are hard to train and deploy in a production environment as they require excessive computational power. for example, the original bert took 4 days to pre-train on 4 tpus. furthermore, there are few limitations of using a non-medical corpus to train a model for medical tasks @xcite . there are medical-specific terms that do not usually exist in general corpora such as news or wikipedia. there are other terms that mean something else in a medical context. the idea of training on domain-specific corpora was explored by @xcite where the authors pre-trained bert on specific domains such as telephone conversations, travel guides, government records and fiction novels. this achieved higher performance for related tasks than the generic version of bert. while versions of bert pre-trained on medical text are publicly available as pointed out in section 2, these models share the same computational power limitations of the original bert. for example, the pre-training of clinicalbert took 18 days on a single gpu. there are no studies we could find as of date that fine-tuned and assessed the performance of distilbert on medical tasks such as ner of phi in medical records. although in the context of deidentification predictions performance is more critical than runtime, the resource limitations may pose a challenge for healthcare organizations to comply with privacy regulations. especially if they need to generate pre-trained embeddings or incrementally fine-tune their models on new data frequently. the question we attempt to answer through this paper is how distilbert performs when finetuned on medical corpora compared to medical pre-trained versions of bert. is it possible to achieve a comparable result to medical pre-trained bert variations such as clinicalbert with a much lighter version such as distilbert? ## method distilbert is based on the concept of knowledge distillation introduced by @xcite . the main characteristic of a machine learning model evaluation is how it performs on unseen data. while high-confidence predictions are picked during inference, there are useful information in low-confidence predictions that can help explain how well a model can generalize. knowledge distillation is a compression algorithm that involves the transfer of such information from the main model, called the teacher, to a smaller distilled version, called the student. further details on knowledge distillation are in the paper by @xcite . distilbert consists of the same two steps as the original bert: pre-training, which in this case creates the student model and fine-tuning which uses the pre-trained student model to train on a custom dataset for a specific task. distilbert was pretrained on the same datasets as the original bert: bookcorpus and wikipedia. the assessment approach was to use the pre-trained distilbert and fine-tune it on i2b2 2010 and i2b2 2014 datasets for ner and compare the results with clinicalbert @xcite and bluebert @xcite that were both pre-trained on medical text. the comparison was done in terms of runtime and f 1 score. the transformers package developed by hugging face co foot_0 was used for all the experiments in this work. its developers are also the creators of distilbert and it hosts a wide variety of pre-trained bert models including the ones mentioned in section 2. the package is implemented in python and this work was implemented in py-torch. throughout this paper, by 'training' we are referring to the supervised learning step that bert and its variants call 'fine-tuning' in order to avoid confusion with hyperparameter tuning. by 'pretraining' we are referring to the unsupervised step that generates the embeddings. the 2014 version has an annotated text of phi labels. the raw data is an xml file with positions of the phi labels. in total, there are 23 different labels with the top 3 accounting to 69% of all label instances (date, doctor, hospital) and the bottom 7 having insignificant counts accounting to near-zero percentages. since it was shown by @xcite that using granular entities for phi achieves better de-identification results than binary classification of whether an entity is a phi, we chose to use all the labels for ner classification instead of binary phi/non-phi classification. access to both datasets was requested through the department of biomedical informatics foot_1 at harvard medical school which is provided for free to researchers and students. the bert model and its variations including distilbert expect ner datasets to be in conll-2003 format introduced by @xcite . it was designed for ner tasks. every line contains the word, a space, and the label of the entity in bio format: b indicates the beginning token of a label, i for inside a multitoken label, and o for a token outside the entities to predict. sequences are separated by two empty lines. in order to produce the training, development and testing datasets in conll format from the raw files, we used the same scripts foot_2 used by clini-calbert authors @xcite . bert variants including distilbert have a hard limit on sequence length set to 512 tokens. some sequences in the raw datasets exceeded that limit. those longer sequences had to be further split to fit the different sequence length experiments. the script referred to in the transformers pack-age's documentation foot_3 was used for splitting longer sequences. the train/test split for both i2b2 2010 and 2014 was already done by i2b2. in order to compare with other papers that used the same datasets as baseline, the train/test split was not modified even though i2b2 2010's training dataset has fewer tokens than its testing dataset. ## results we can draw the following insights from the results presented in ## https://github.com/google-research/ bert another aspect of the results we were interested in was runtime. as mentioned in section 3, the original bert is heavy to use requiring significant computing resources. distilbert runtime was 43% faster on average than the medical variants of bert. it also produced a model that was consistently 60% smaller in size than bluebert and clinicalbert. the per-label performance for all models and both datasets is shown in appendix a. we can draw the following insights from the per-label comparison. as shown by support numbers column, the testing dataset for i2b2 2014 did not have significant counts for entities such as fax, device and email therefore producing unpredictable f 1 scores. this subsequently drove the average f 1 lower. however, of the top 4 frequent labels (date, doctor, patient and hospital), date had the highest f 1 score. on the other hand, hospi-tal had the lowest f 1 with significant support number (877 instances) achieving 48 f 1 score which contributed to a lower micro f 1 average for distilbert uncased. on the other hand, distil-bert cased model performed significantly better for the hospital label achieving 88 f 1 score. as discussed earlier, casing features are important in the context of ner tasks. for english nouns, casing is particularly important. for example, hospital names are written with a capital first letter. for i2b2 2010, since labels are all medical concepts, distilbert had trouble recognizing all 3 entities of treatment, problem and test achieving an f 1 ranging from 78 to 82. for comparison, bluebert achieved 84-85 for all 3 entities and clinicalbert achieved 87. the parameters that yielded the best performance out of all 40 runs are shown in in terms of relative performance, distilbert model's f 1 score was, on average, 95% that of medically pre-trained bert score for i2b2 2010 containing medical terms but on par for i2b2 2014 dataset. this result is 2 points lower than reported by @xcite for question answering task. the performance degradation of using a distilled model is therefore task-as well as data-specific. ## conclusion in this work the main contribution was a full performance assessment of distilbert in terms of