ReMAG-KR: Retrieval and Medically Assisted Generation with Knowledge Reduction for Medical Question Answering (2024)

Abstract: 001 Large Language Models (LLMs) have significant 002 potential for facilitating intelligent end-user ap003 plications in healthcare. However, hallucinations 004 remain an inherent problem with LLMs, mak005 ing it crucial to address this issue with extensive 006 medical knowledge and data. In this work, we 007 propose a Retrieve-and-Medically-Augmented008 Generation with Knowledge Reduction (ReMAG009 KR) pipeline, employing a carefully curated 010 knowledge base using cross-encoder re-ranking 011 strategies. The pipeline is tested on medical 012 MCQ-based QA datasets as well as general QA 013 datasets. It was observed that when the knowl014 edge base is reduced, the model’s performance 015 decreases by 2-8%, while the inference time im016 proves by 47%. 017

text: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 4: Student Research Workshop), pages 140–145 August 11-16, 2024 ©2024 Association for Computational Linguistics
ReMAG-KR: Retrieval and Medically Assisted Generation with Knowledge Reduction for Medical Question Answering
Anonymous ACL submission
Abstract001
Large Language Models (LLMs) have significant002 potential for facilitating intelligent end-user ap-003 plications in healthcare. However, hallucinations004 remain an inherent problem with LLMs, mak-005 ing it crucial to address this issue with extensive006 medical knowledge and data. In this work, we007 propose a Retrieve-and-Medically-Augmented-008 Generation with Knowledge Reduction (ReMAG-009 KR) pipeline, employing a carefully curated010 knowledge base using cross-encoder re-ranking011 strategies. The pipeline is tested on medical012 MCQ-based QA datasets as well as general QA013 datasets. It was observed that when the knowl-014 edge base is reduced, the model’s performance015 decreases by 2-8%, while the inference time im-016 proves by 47%.017
1 Introduction018
Large Language Models (LLMs) like GPT-4019 (Achiam et al., 2023), LLaMA-2 (Touvron et al.,020 2023a), and LLaMA-3 (Touvron et al., 2023b) have021 become highly efficient text generation tools with022 a significant variety of potential applications in a023 wide range of domains, like business, education, and024 healthcare. In healthcare, the potential to transform025 challenging tasks such as patient education (Jin et al.,026 2024), report generation (Shoham and Rappoport,027 2023), and drug discovery (Kormilitzin et al., 2021;028 Unnikrishnan et al., 2023) is exemplary. This is pri-029 marily due to their ability to analyze large amounts030 of textual data and generate high-quality meaning-031 ful text as per end-user task requirements. However,032 there are specific challenges with deploying them033 in the healthcare industry. The shortage of medical034 data available for training/consumption by LLMs is035 one of the primary reasons for concern. A critical036 hurdle is the propensity of LLMs to produce false037
medical information (often termed hallucinations), 038 misleading both patients and medical professionals. 039 Additionally, in case of instructions that are too ex- 040 plicit or devoid of important details, LLMs fail to 041 produce optimal results, which reduces their effi- 042 cacy. LLMs may also reinforce biases learned from 043 the training data, producing biased results towards 044 particular groups of people based on constructs like 045 gender, ethnicity, and socio-economic status. 046
For general tasks, the application of the concept of 047 Retrieval-Augmented Generation (RAG) has shown 048 promise. RAG systems incorporate external infor- 049 mation retrieval into the LLM architecture. Previ- 050 ous research, such as Almanac (Zakka et al., 2024) 051 and ChatENT (Long et al., 2023), has demonstrated 052 improved LLM accuracy and reliability with this 053 method. However, this kind of integration may also 054 include unrelated or incorrect information, which 055 could undermine the legitimacy and efficacy of the 056 LLM. Including external knowledge sources raises 057 issues with data consistency, privacy, security, and 058 legal consequences. Furthermore, these methods 059 frequently call for indexing and storing massive 060 datasets, sometimes surpassing 200GB. Although 061 RAG approaches perform excellently in general 062 question-answering tasks, there is still uncertainty 063 about their efficiency in healthcare. Regarding effi- 064 ciency, retrievers trained on generic data often fall 065 short of those optimized for particular domains (Li 066 et al., 2022). This emphasizes the need for domain- 067 specific training data, which can be costly and time- 068 consuming to create, particularly in specialized fields 069 like medicine. Moreover, conventional RAG tech- 070 niques train the LLM and retriever separately (Stein- 071 berg et al., 2021; Agrawal et al., 2022), while other 072 approaches include joint training of retrievers and 073 LLMs (Wang et al., 2024). The retrieved informa- 074
1
140
tion and the LLM’s capacity to process it for accurate075 output may generally lack semantic depth due to the076 nature of the training (Sarthi et al., 2024).077 To address these challenges without additional078 computational costs, we propose a systematic ap-079 proach that integrates RAG models built on a care-080 fully curated knowledge base, with support for cross-081 encoder re-ranking strategies. First, keywords and082 entities from each query or question are extracted.083 Then, a web crawl is conducted to find each entity’s084 top-15 relevant documents, which are used to build085 the knowledge base. Next, the retrieval based on the086 query and re-ranking of the results using MedCPT087 (Jin et al., 2023) is performed. Finally, responses088 are generated using LLMs, specifically LLaMa2 and089 LLaMa3. The rest of the article is structured as fol-090 lows. Section 2 presents a detailed discussion on the091 proposed approach. Section 3 presents a discussion092 on the experiments performed and results observed,093 followed by conclusion and future work.094
2 Methodology095
Fig. 1 depicts the proposed approach consisting of096 five key phases – Keyword extraction, Document Re-097 trieval, Knowledge base construction, Cross-encoder098 re-ranking, and response generation using LLMs.099 Given a set of medical questions Q, a set of key-100 words KQ are extracted using KeyBERT (Grooten-101 dorst, 2020). For each keyword k ∈ KQ, a ranked102 set of 15 relevant documents di are retrieved from103 the PubMed database, resulting in a comprehensive104 collection of medical documents D∗ containing perti-105 nent medical knowledge. Formally, for each question106 qi ∈ Q, there exists a corresponding correct answer107 a∗i within a set of options Ai, such that a ∗ i ∈ Ai.108 The model M utilizes the query q and relevant doc-109 ument d to produce a predicted answer (as per Eq.110 (1), where, d ∈ D∗ and dR is the retrieved document111 based on the query, and Eq. (2) where θ represents112 the model’s parameters).113
Document Retrieval dR = p(d | q) (1)114 115 Answer Prediction a = p(a | q, dR, θ) (2)116 In the Keyword Extraction phase, at least three117 keywords are extracted from all queries. For each118 query Q, KeyBERT is used to extract at least three119 related medical keywords or key phrases KQ. This120
ensures that individual words and phrases relevant 121 to the medical context are captured accurately. Fol- 122 lowing this, in the Knowledge Base Indexing and 123 Storage phase, the PubMed API is employed to re- 124 trieve 15 relevant articles for each identified keyword 125 or keyphrase. This retrieval process results in a sub- 126 stantial corpus of about 600,000 articles, providing 127 a focused subset of the extensive PubMed database 128 (Canese and Weis, 2013), consisting of 24.9 million 129 articles. The collected articles are then transformed 130 into embedding vectors through the BAAI embed- 131
2
ding model (Zhang et al., 2024). This model con-132 verts text data into a format that can be efficiently133 processed for similarity searches. Finally, the FAISS134 (Facebook AI Similarity Search) VectorStore index135 is used to store these embedding vectors. FAISS136 is optimized for high-speed similarity searches on137 large datasets, making it suitable for handling the138 extensive medical corpus generated.139 For facilitating retrieval, MedCPT was used in the140 process of document retrieval. For this, we compute141 the cosine similarity score between the query embed-142 ding q and each document embedding di. The top143 k documents were chosen based on query similar-144 ity. MedCPT was utilized to streamline the retrieval145 process and ensure that the most relevant documents146 were retrieved. For this, the cosine similarity score147 was computed between the query embedding q and148 the embedding of each document di. The cos-sim149 score measures how similar two vectors are in orien-150 tation and magnitude, with a higher score indicating151 greater similarity. By calculating this score for each152 document in the database, the system can effectively153 rank them based on their relevance to the query. Us-154 ing the computed scores, the top k documents were155 selected for retrieval, ensuring that the retrieved doc-156 uments are closely aligned with the user’s query.157 The cross-encoder re-ranker MedCPT is advanta-158 geous in re-ranking the top k extracted articles for159 generating the top n articles (where n = k), enhanc-160 ing the relevance of the information produced. Med-161 CPT was chosen as retriever and re-ranker due to its162 First-stage dense retriever (MedCPT retriever) and163 the Second-stage re-ranker (MedCPT re-ranker). The164 MedCPT retriever contains a query encoder (QEnc)165 and an article encoder (DEnc), both initialized by166 PubMedBERT. It is trained on 255M query-article167 pairs from PubMed search logs and in-batch nega-168 tives. The MedCPT re-ranker is a transformer cross-169 encoder (CrossEnc) initialized by PubMedBERT. It170 is trained on 18M semantic query-article pairs and171 localized negatives derived from the pre-trained Med-172 CPT retriever.173 Upon re-ranking retrieved articles, they are174 merged with the original query and provided as input175 to the LLM, which generates a response, represented176 as a, based on the amalgamation of the query and177 the re-ranked articles. The generated response is178 then evaluated by comparing it with the ground-truth179 answers, serving as a metric for assessing the perfor- 180 mance of the LLM in understanding and responding 181 to the given query. We have utilized two specific 182 LLMs for our experiments, namely LLaMA-2 and 183 LLaMA-3. These models have been selected based 184 on their capabilities and suitability for the task at 185 hand. We aim to evaluate these LLMs’ effectiveness 186 in generating accurate responses when presented 187 with queries and relevant document contexts. 188 3 Experiments and Results 189 Experiments were conducted on the benchmark MI- 190 RAGE dataset (Xiong et al., 2024) for the mul- 191 tiple choice questions-based QA tasks. This in- 192 cluded 7,663 questions from five commonly used QA 193 datasets in biomedicine (MMLU-Med, MedQA-US, 194 MedMCQA, PubMedQA (Jin et al., 2019), BioASQ 195 (Y/N)) (Tsatsaronis et al., 2015). For Subjective QA 196 task, the datasets LiveQA (Abacha et al., 2017) and 197 ExpertQA-Med (Malaviya et al., 2023) were cho- 198 sen, with 3,479 subjective questions and answers. 199 Standard metrics like accuracy, precision, recall, and 200 F1-score were used for the evaluation. The gener- 201 ated text quality and relevancy were assessed using 202 BLEURT, BERTScore, MoverScore, and ROUGE- 203 L. For MCQ-based QA tasks, the MED-RAG model 204 was used as the baseline, while KG-Rank (Yang et al., 205 2024) was considered for subjective tasks due to its 206 novelty and outstanding scores. 207 3.1 Results and Discussion 208 MCQ-based tasks: Table 1 shows the results for 209 this task, and it is evident that the proposed ReMAG- 210 KR underperformed on datasets including MMLU- 211 Med, MedQA-US, MedMCQA, PubMedQA, and 212 BioASQ-Y/N. Compared to MEDRAG’s 73.09 av- 213 erage accuracy, our approach produced 66.32. Like- 214 wise, our approach averaged 58.74 for the F1 score, 215 whereas MEDRAG scored 66.69. Despite the lag 216 in performance, the proposed ReMAG-KR showed 217 a notable efficiency advantage, as seen in Table 218 3. The inference time was nearly one-third that 219 of MEDRAG, primarily because our knowledge 220 base contains only 600,000 documents compared 221 to MEDRAG’s extensive 25 million corpus. 222 Subjective tasks: Using LLaMA-2 and LLaMA- 223 3 models, we analyzed the ExpertQA-Med and 224
3
LiveQA datasets and compared the effectiveness of225 the RR and ReMAG-KR approaches (Refer Table226 2. With LLaMA-2 and LLaMA-3, respectively, KG-227 Rank obtained a ROUGE-L of 28.02 and 29.03 and a228 BERTScore of 86.01 and 86.93 for ExpertQA-Med.229 ReMAG-KR obtained BERTScore and ROUGE-L230 scores of 82.43, 84.23, 28.32, and 28.72, respec-231 tively. For LiveQA, KG-Rank obtained BERTScores232 of 82.02 and 83.38 in addition to ROUGE-L scores233 of 19.72 and 20.26. For BERTScore, ReMAG-KR234 scored 79.77 and 81.97, and for ROUGE-L, 17.84235 and 18.84. Both techniques showed comparable per-236 formance with equal inference times.237
4 Conclusion and Future Work 238
An approach for LLM-based retrieval and medically 239 assisted generation with a tactically reduced knowl- 240 edge base was presented in this article. Experiments 241 revealed that our approach reduces inference time 242 by 47% with a small compromise in performance 243 (of around 2-8%). Performance for subjective QA 244 tasks was also comparable with the state-of-the-art 245 approaches in this field. We plan to extend the pro- 246 posed approach by enriching the quality of retrieved 247 documents, while maintaining a reduced inference 248 time. Simple, vanilla RAG-based approaches fail to 249 capture semantically deep information hidden within 250 medical text, thus, the use of a single corpus for 251 generating a knowledge base encompassing multi- 252 ple sources of information could be attempted. We 253 also plan to introduce two other components into 254 the RAG pipeline – multi-hop question answering 255 and question decomposition. This involves breaking 256 down a complex query into sub-queries and enrich- 257 ing the quality of retrieved documents. Adopting 258 domain-specific models like PMC and MedLLaMA 259 may further boost the model’s ability to handle the 260 intricacies and nuances inherent in medical data. 261
4
References262 Asma Ben Abacha, Eugene Agichtein, Yuval Pinter, and263 Dina Demner-Fushman. 2017. Overview of the med-264 ical question answering task at trec 2017 liveqa. In265 TREC, pages 1–12.266
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama267 Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo268 Almeida, Janko Altenschmidt, Sam Altman, Shyamal269 Anadkat, et al. 2023. GPT-4 technical report. arXiv270 preprint arXiv:2303.08774.271
Monica Agrawal, Stefan Hegselmann, Hunter Lang, Yoon272 Kim, and David Sontag. 2022. Large language models273 are few-shot clinical information extractors. arXiv274 preprint arXiv:2205.12689.275
Kathi Canese and Sarah Weis. 2013. Pubmed: the biblio-276 graphic database. The NCBI handbook, 2(1).277
Maarten Grootendorst. 2020. Keybert: Minimal key-278 word extraction with bert. https://github.com/279 MaartenGr/KeyBERT.280
Mingyu Jin, Qinkai Yu, Chong Zhang, Dong Shu, Suiyuan281 Zhu, Mengnan Du, Yongfeng Zhang, and Yanda282 Meng. 2024. Health-llm: Personalized retrieval-283 augmented disease prediction model. arXiv preprint284 arXiv:2402.00746.285
Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William W286 Cohen, and Xinghua Lu. 2019. Pubmedqa: A dataset287 for biomedical research question answering. arXiv288 preprint arXiv:1909.06146.289
Qiao Jin, Won Kim, Qingyu Chen, Donald C Comeau,290 Lana Yeganova, W John Wilbur, and Zhiyong Lu.291 2023. Medcpt: Contrastive pre-trained transform-292 ers with large-scale pubmed search logs for zero-293 shot biomedical information retrieval. Bioinformatics,294 39(11):btad651.295
Andrey Kormilitzin, Nemanja Vaci, Qiang Liu, and Alejo296 Nevado-Holgado. 2021. Med7: A transferable clini-297 cal natural language processing model for electronic298 health records. Artificial Intelligence in Medicine,299 118:102086.300
Huayang Li, Yixuan Su, Deng Cai, Yan Wang, and Lemao301 Liu. 2022. A survey on retrieval-augmented text gener-302 ation. arXiv preprint arXiv:2202.01110.303
Cai Long, Deepak Subburam, Kayle Lowe, André dos304 Santos, Jessica Zhang, Sang Hwang, Neil Saduka, Yoav305 Horev, Tao Su, David Cote, et al. 2023. Chatent:306 Augmented large language model for expert knowl-307 edge retrieval in otolaryngology-head and neck surgery.308 medRxiv, pages 2023–08.309
Chaitanya Malaviya, Subin Lee, Sihao Chen, Elizabeth 310 Sieber, Mark Yatskar, and Dan Roth. 2023. Expertqa: 311 Expert-curated questions and attributed answers. arXiv 312 preprint arXiv:2309.07852. 313
Parth Sarthi, Salman Abdullah, Aditi Tuli, Shubh Khanna, 314 Anna Goldie, and Christopher D Manning. 2024. Rap- 315 tor: Recursive abstractive processing for tree-organized 316 retrieval. arXiv preprint arXiv:2401.18059. 317
Ofir Ben Shoham and Nadav Rappoport. 2023. Cpllm: 318 Clinical prediction with large language models. arXiv 319 preprint arXiv:2309.11295. 320
Ethan Steinberg, Ken Jung, Jason A Fries, Conor K 321 Corbin, Stephen R Pfohl, and Nigam H Shah. 2021. 322 Language models are an effective representation learn- 323 ing technique for electronic health record data. Journal 324 of biomedical informatics, 113:103637. 325
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier 326 Martinet, Marie-Anne Lachaux, Timothée Lacroix, 327 Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal 328 Azhar, et al. 2023a. Llama: Open and effi- 329 cient foundation language models. arXiv preprint 330 arXiv:2302.13971. 331
Hugo Touvron, Louis Martin, Kevin Stone, Peter Al- 332 bert, Amjad Almahairi, Yasmine Babaei, Nikolay Bash- 333 lykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhos- 334 ale, et al. 2023b. Llama 2: Open foundation and fine- 335 tuned chat models. arXiv preprint arXiv:2307.09288. 336
George Tsatsaronis, Georgios Balikas, Prodromos 337 Malakasiotis, Ioannis Partalas, Matthias Zschunke, 338 Michael R Alvers, Dirk Weissenborn, Anastasia 339 Krithara, Sergios Petridis, Dimitris Polychronopou- 340 los, et al. 2015. An overview of the bioasq large-scale 341 biomedical semantic indexing and question answering 342 competition. BMC bioinformatics, 16:1–28. 343
Reshma Unnikrishnan, Sowmya Kamath, and VS Anan- 344 thanarayana. 2023. Efficient parameter tuning of neu- 345 ral foundation models for drug perspective prediction 346 from unstructured socio-medical data. Engineering 347 Applications of Artificial Intelligence, 123:106214. 348
Junda Wang, Zhichao Yang, Zonghai Yao, and Hong Yu. 349 2024. Jmlr: Joint medical llm and retrieval training for 350 enhancing reasoning and professional question answer- 351 ing capability. arXiv preprint arXiv:2402.17887. 352
Guangzhi Xiong, Qiao Jin, Zhiyong Lu, and 353 Aidong Zhang. 2024. Benchmarking retrieval- 354 augmented generation for medicine. arXiv preprint 355 arXiv:2402.13178. 356
Rui Yang, Haoran Liu, Qingcheng Zeng, Yu He Ke, 357 Wanxin Li, Lechao Cheng, Qingyu Chen, James Caver- 358 lee, Yutaka Matsuo, and Irene Li. 2024. Kg-rank: En- 359 hancing large language models for medical qa with 360
5
knowledge graphs and ranking techniques. arXiv361 preprint arXiv:2403.05881.362
Cyril Zakka, Rohan Shad, Akash Chaurasia, Alex R Dalal,363 Jennifer L Kim, Michael Moor, Robyn Fong, Curran364 Phillips, Kevin Alexander, Euan Ashley, et al. 2024.365 Almanac—retrieval-augmented language models for366 clinical medicine. NEJM AI, 1(2):AIoa2300068.367
Mingtian Zhang, Shawn Lan, Peter Hayes, and David Bar-368 ber. 2024. Mafin: Enhancing black-box embeddings369 with model augmented fine-tuning. arXiv preprint370 arXiv:2402.12177.371
6

heading: Future Work