REMED: Retrieval-Augmented Medical Document Query Responding with Embedding Fine-Tuning

Tianqi Pang, Kehui Tan, Yujun Yao, Xiangyang Liu, Fanlong Meng, Chenyou Fan, Xiaofan Zhang

Abstract: While advanced Large Language Models (LLMs) exhibit considerable promise, their tendency to generate unreliable information poses significant challenges, particularly in high-risk domains like healthcare. However, the advent of Retrieval-Augmented Generation (RAG) offers a novel solution tailored for the medical realm. This study further enhances retrieval accuracy by introducing REMED, a specialized medical document retrieval framework designed to address the hallucination problem prevalent in LLMs. The REMED framework integrates dataset construction, an efficient embedding fine-tuning EM-FT model, retrieval-augmented generation, and human evaluation of LLM responses. The EM-FT model can end-to-end fine-tune the medical sentence representations in large pre-trained models through an efficient embedding fine-tuning method, thereby enhancing the performance of medical retrieval. We adopt contrastive learning as the loss function to optimize the performance of the EM-FT model, enabling it to accurately capture the similarity between query and relevant documents. This approach not only improves the retrieval accuracy of positively related contents but also effectively reduces the matching with negatively related contents. Compared to direct dense vector retrieval, fine-tuning query and content vectors first and then performing dense retrieval tasks significantly improved the performance. Through validation on two datasets, we demonstrate that our EM-FT method improves recall and precision on MMD by 3.2%-6.0% and on MPD by 14.4%-42.6% compared to using the embedding model directly for retrieval. Furthermore, through human evaluation on the PULSE-7Bv5 model, we further confirm the effectiveness of our retrieval results in improving the quality of generated text.