Tools and Technology Seminar 11/07/2024 - Yufeng Zhang

Hi everyone. The title for today's presentation will be 'Using the Large Language Model for the Radiology Report,' but this is part of a huge project called Necklace, and it involves a lot of background. I'll start from Necklace and then switch to the main method I propose and use in this research. Necklace is an abbreviation for the project: Necrotizing Enterocolitis Labeling and Computer-Enhanced Diagnosis. Today's presentation will be divided into four parts: background of the disease and why we wanted to build this system, what the challenges are, and there are two sub-aims. The first aim is what we are going to talk about today, and the second one will be for the future. NEC is a very serious gastrointestinal disease in premature infants that causes inflammation and tissue death in the intestine. The disease can progress really fast. It affects around 7% of infants whose birth weight is between 500 and 1,500 grams, and the mortality is alarmingly high, ranging from 20% to 30%. We can see from the images shown here: on the left side is the abdominal X-ray, and this is the abnormal structure where the bowel or intestine is full of air. On the right side, the normal abdominal X-ray is very clear. The current data we have is text and image pair data: the text is the Radiology report, and the images are the abdominal X-ray. In the Radiology reports, there are several sections: the title, the time of the examination performed, followed by clinical history, comparison, impression, and routine tests. The impression section mainly describes what the radiologist observed in the radiology images. We have around 15,000 data points in total. We discussed with clinicians, and based on their needs, they want the system to perform such that when the radiology images come in, the model immediately flags them as abnormal or normal. So the system would be purely based on the images, but the key problem for us is that we need the labels extracted from the Radiology report, which is unstructured text. How should we extract labels from this unstructured data? There are two main challenges. The first is the clinical aspect: timely diagnosis is very crucial because the disease can progress really fast, so any delay in review can postpone necessary treatment and worsen the patient's condition. In addition, this is highly reliant on expertise; not all general physicians can make conclusions from the radiology images; it requires very specific domain knowledge. From the technical aspect, as I mentioned, we don't have labels. The label is implicitly stated in the impression section of the Radiology report, so how should we extract labels from the report? The second technical challenge is that there is no public available abdominal X-ray dataset. We only have 15,000 pairs of images and text. Based on other work on chest X-ray or other medical imaging, there are a lot of public datasets they can use, so they have a really good initialization and a good understanding of what the normal abdominal X-ray should be, but we don't have such knowledge right now. That leads to our Aim 1: 'Human Expert-Facilitated, Privacy-Preserving Large Language Model for Abdominal X-ray Report Labeling.' We'll talk about this part today. After we build the whole system, on the left side is how the whole Necklace model is built. The first box is what we are going to talk about today: we have the Radiology reports, and we invited human experts to annotate these Radiology reports to tell whether there is an NEC description in the report or not. We then use this annotated dataset to train a model, and we use the trained model to label the remaining unannotated dataset. Then the second box is to train an image classification model, which will be mainly used in the whole clinical support system. Tremendous research has been mainly focused on chest X-ray, and there are a lot of chest X-ray public datasets out there, like MIMIC and ChestX-ray14. There are millions of chest X-ray public datasets out there. Benefiting from this extensive data, there are numerous and rapidly evolving chest X-ray labeling tools developed. But currently, we only have a single facility abdominal X-ray dataset, so we have no available public tools and we don't have enough datasets. But we can still learn from how people build these chest X-ray tools. Let's do a quick literature review of how people in the past built chest X-ray related tools so we can build our own abdominal X-ray related tools. The first tool developed for chest X-ray report labeling is called NegBio or CheXpert. They are solely based on regular expressions on the surface text, and they are very limited when attempting to capture complex syntactic constructions such as long noun phrases. The limitation is obvious: they usually miss the nuanced and varied language patterns in Radiology reports. Sometimes the description of certain diseases or observations in the Radiology report can be very ambiguous, so the regular expressions don't have the capability to capture such nuanced expressions. In addition, they need to identify negation in the sentence, like 'no evidence of.' They are very prone to typos and ambiguities, and they are out of date because the papers were published around 2018 and 2019. After Transformers were introduced in this field since 2020 and 2021, people kind of switched to BERT-based models like CheXbert or CheXbert Plus. First, they used the previous rule-based labelers to annotate the unannotated Radiology reports and selected around 1,000 reports and invited some board-certified radiologists to label them. So the first step is an initial annotation result. This method trains the model using very initial low-quality annotations and then fine-tunes using the expert-curated dataset to improve model performance. The improvement is that it utilizes the direct report-to-label mapping from clinicians, and it balances between rule-based labelers and manual annotations because manual annotations require a lot of clinician time and effort, so it reduces cost and labeling effort while maintaining good model accuracy. As LLMs were introduced last year, a lot of people are exploring whether they can be utilized in this field, but most of the work uses in-context learning. In-context learning means we give some examples to the GPT models and give them an unannotated report and ask them to annotate it based on the several examples we feed into the model. The core is to develop a finely crafted prompt that can appropriately describe the whole problem and help GPT understand what you really want to do. For us, it's impossible because GPT is a closed-source tool, and if we wanted to use GPT, we would need to send our data to OpenAI, which is very sensitive patient data and would compromise patient privacy. So what should we do? This is the performance comparison of all previous methods, and it turns out that the GPT model, the large language model, shows really good model performance compared to the standard one and the BERT-based models. So it shows promise, but we cannot use GPT. So let's talk about the task here. NEC is air in the intestine, but it has three subtypes. The first is called pneumatosis, which is air in the intestinal wall. The second is PVG, portal venous gas, meaning the air flows along the vein into the liver. The third is called free air, meaning the air is all over the abdomen, which is the most severe. We invited clinicians to annotate the Radiology report based on the impression section and give three labels for the first task: whether there is air in the abdomen. The second is whether it is pneumatosis, PVG, or free air. Sometimes because the impression says 'bubbly lucencies' or something, they cannot tell if it's stool or air, so they label it as uncertain. So it's a multi-class classification problem: positive, negative, uncertain. The same for the next sub-type classification task. This is the method we propose. We have around 15,000 samples in the dataset, and we selected 2,500 samples and invited three pediatricians who have domain knowledge to help us label this dataset. In addition, we used 1,600 of them for training purposes and the rest for hyperparameter tuning and testing. We designed instruction tuning prompts that describe this whole task and trained large language models: Gemma and Mistral AI models. They are open-source models that can be trained and deployed locally. After we well-trained this model and achieved really good model performance, we wanted to do model distillation to transfer the knowledge from the LLM to the BERT model because BERT is a relatively small model and can quickly do inference. This is the whole flowchart. There are two questions we addressed using this framework. The first is why we use LLM: because there is no public dataset, and the LLM in the previous literature review can be considered an expert model; in addition, it reduces human labeling cost. The second is why we want to do distillation: first, to improve model latency because using LLM takes a lot of time; second, to reduce cost because LLM requires a lot of computing resources; and also for privacy, as I mentioned, BERT only stores the weights, but the LLM stores not only the weights but also the token relationships. This is the prompt I designed. First, I told the model this is a pediatric Radiology report, and the input would be the impression section from the Radiology report. There are four multiple-choice questions. The first would be: 'Does the child have NEC? A: Yes, B: No, C: Uncertain.' Very similar for the other subtypes. To facilitate downstream post-processing, I asked the model to give the answer using A, B, or C in an answer tag so I can easily extract the answers from the output and evaluate model performance. During training, we know the label, so we let the LLM generate the answer themselves. These are the two models we used: Gemma and Mistral. They both use multi-query attention, group attention, sliding window attention, which reduce the parameters trained in the self-attention and accelerate training and inference time and reduce computing resources. This is the configuration I used: I used a single Nvidia A40 GPU and requested 16 gigabytes. Usually, large language models require hundreds of gigabytes, but the reason I only requested 16 gigabytes is because I used a method called QLoRA: quantization plus low-rank adaptation. These are the hyperparameters I used for model training. I made a table here. The epoch is actually 10, so I only used 10 epochs and achieved really good model performance compared to the baseline. QLoRA is: in the original model, we do full fine-tuning, but for large language models, it's impossible for our lab to do full fine-tuning, so we used QLoRA. LoRA introduces auxiliary matrices that only learn the update of the weight instead of the whole weight. Q means quantization: the original model uses float16, and we can first quantize the model into a 4-bit frozen base model and only learn several adapters during training, not the whole layers. These adapters are only applied to several important layers like in self-attention: the query, key, and value matrices. The time complexity for these matrices can be O(n^2), but if we use the adapter, we can reduce the time complexity a lot. During inference, we do dequantization: the 4-bit frozen model is converted back to the original float16 model, and we add the output from the adapters. Instruction fine-tuning... I don't think this is important. The last proposed component is model distillation. There are three main model distillation paradigms. The first is online distillation: we have two models, a teacher model and a student model, and they are trying to learn the patterns of the teacher model, but the teacher model is also updating at the same time. Usually, the loss design here is KL divergence, trying to match the two distributions. The second is offline distillation, which is what we used: we first well-train a teacher model and use the teacher model to annotate a distilled dataset, and then we use the distilled dataset to train a student model. The teacher model is fixed and no longer updated, but the teacher model will train the data and use the data to tell how the student model should behave. The third is called parametric knowledge transfer, which only applies to large language models. In their paper, they have a lot of LoRA layers and use SVD to choose which layer is the most important and only apply those important LoRA layers to the student model. So that's why it's called parametric knowledge transfer. We are going to use the second one because our teacher model achieved really good model performance, so we don't need to update the teacher model again. This is how the data is split. We have 2,500 manually annotated reports, and we split them into training, testing, and development. Observe the data distribution table: for the positive cases and uncertain cases, the number of these samples is relatively small compared to the negative samples. So it's a super imbalanced dataset and not very large. In previous literature, when people train a BERT model, they use millions of data points, but we only have a small amount. These are the results comparing LLMs against fine-tuned BERT models. The BERT models we used are pre-trained: one is called ClinicalBERT, trained on the MIMIC dataset, and the other is called BlueBERT, pre-trained on PubMed abstracts. That can explain why BlueBERT is slightly better than ClinicalBERT because MIMIC is for patients aged over 18, but we are focusing on NEC patients who are premature infants, so ClinicalBERT doesn't perform well because there is not enough relevant knowledge in the pre-trained model. These two BERT-based models don't perform well: the F1 score for positive and uncertain is about 0.26 and 0.38. But if we use Gemma 2 billion, Gemma 7 billion, or Mistral 7 billion, the model performance is actually really high. For NEC positive cases, the F1 score can be up to 0.9, and for uncertain cases, it can be up to 0.778. We ran the experiment three times, and the standard deviation is relatively small. So that is how the LLM performs against the BERT models. This time we are focusing on the Gemma 2B and Gemma 7B models. We can conclude from Gemma 2B and Gemma 7B: the ratio here means the proportion of the training dataset used during fine-tuning. 0.2 means we used only 20% of the training dataset, so we utilized around 300 training samples. For Gemma 2B, we can see that as the fine-tuning size increases, the model performance actually increases as well. The same story for Gemma 7B and Mistral 7B. We will not talk about Mistral 7B here because it doesn't perform very stably and the standard deviation is really high. So we focus on Gemma 2B and Gemma 7B. Compared to Gemma 2B and Gemma 7B, no matter what the fine-tuning ratio is, the Gemma 7B model performance is consistently better than the Gemma 2B model. So we can conclude that a larger model has better capability here. The second conclusion is that when we increase fine-tuning size, the model performance actually increases. That is not very surprising because we have more data to train the model, and the model can better model the relationship between input and output. But we actually observe a jump when we increase the fine-tuning dataset from 20% to 40%. I initially thought that would be something like emergent ability, but emergent ability is only talking about when the model size increases to a certain size, they can solve a lot of complex tasks they could never solve before. I'm not so sure if this also works for the fine-tuning dataset. I only did empirical experiments, so I didn't study the theory behind. It's just an observation: as the fine-tuning data increases, model performance increases. There is a jump here for the Gemma 2B model, but for the Gemma 7B model, because it is a relatively large model and already has a lot of pre-trained knowledge, even when the fine-tuning size is 0.2, the model performance remains relatively high for positive and uncertain. So the next conclusion is that smaller models benefit more from larger datasets. The last one is that for PVG and free air, they only have like four to five samples in the dataset, so even for the Gemma 7B model, they need to have a full fine-tuning dataset to achieve a 0.9 or 0.8 F1 score. 0.2 is not enough for them to perform well. The last part is how the distilled model works. In the previous BERT models, no matter if it was ClinicalBERT or BlueBERT, the F1 score for NEC or pneumatosis was very low. We wanted to know if we use LLM to teach the BERT model whether it can perform better compared to its baselines. We can see that after we did model distillation, Gemma 2B, 7B, and Mistral 7B all increased model performance no matter what teacher model they used, and it is on par with the LLM teacher model itself. So we have two conclusions: this student model can achieve comparable model performance to the LLM teacher, and a better teacher model results in a better student model. For example, Gemma 7B is doing better than Gemma 2B as a teacher model. The takeaway messages for LLM in abdominal X-ray classification are: first, different LLMs show varying capability in labeling abdominal X-ray reports, and larger models consistently outperform smaller ones. Second, increasing the size of the fine-tuning dataset significantly boosts model performance, particularly for smaller models and classes with limited examples. Third, using LLM as a teacher model can substantially improve the performance of the BERT model, making it comparable to LLMs. That concludes all the previous results. There are future discussions. I said why we choose Gemma and Mistral instead of GPT: because Gemma and Mistral can be run locally, preventing sensitive data leakage, and distilled models can be easily fine-tuned and transmitted because they only store the weights. The limitation is that we simplified the model into a classification problem. This is what clinicians requested because they wanted the images to come in and the label to come out. But there are a lot of people out there doing report summary or report generation. That will probably be the future of this project because we will never know what classification we want in the future. In addition, this is retrospective data and only from a single facility, so that will introduce bias into the model; the model may not generalize well to other datasets. Now we are collaborating with the NEC Society and hospitals to get more data. To address the previous limitations, we will probably improve model interpretability, for example, introduce chain-of-thought explanations in the loop to improve reasoning abilities using direct preference optimization. Also, for multicenter medical data, because we have a lot of collaborations and we are not sure if they want to send data to us, we can use federated learning to preserve their local privacy in this process. I think that will be all I wanted to talk about today because the rest is the image part, which is probably irrelevant to the topic. Are there any questions?

Interesting talk. I see that your comparisons are for different LLMs and different language BERT models, right? Different language, you compare different models, different parameters. Even though we are comparing these, there will definitely be mislabeled cases and other things. Were you able to compare it with the actual label data from the clinicians? How well do all of these language models perform compared to that?

This is the results compared to the clinicians. Actually, we split the data into training, validation, and testing, and the results I reported here are the model performance on the testing dataset.

Are there any other metrics other than the precision and recall that can be used to generate... Are you talking about what else input?

Currently, only the Radiology report is used. Actually, the impression section, like 'the gas extending to the rectum,' only that section is used during training. The average token number is around...