JMLR: Joint Medical LLM and Retrieval Training for Enhancing Reasoning and Professional Question Answering Capability (2402.17887v4)

Abstract: LLMs have demonstrated a remarkable potential in medical knowledge acquisition and question-answering. However, LLMs can potentially hallucinate and yield factually incorrect outcomes, even with domain-specific pretraining. Previously, retrieval augmented generation (RAG) has limited success in addressing hallucinations. Unlike previous methods in RAG where the retrieval model was trained separately from the LLM, we introduce JMLR (for Jointly trains LLM and information Retrieval) during the fine-tuning phase. The synchronized training mechanism enhances JMLR's ability to retrieve clinical guidelines and leverage medical knowledge to reason and answer questions and reduces the demand for computational resources. We evaluated JMLR on the important medical question-answering application. Our experimental results demonstrate that JMLR-13B (70.5%) outperforms a previous state-of-the-art open-source model using conventional pre-training and fine-tuning Meditron-70B (68.9%) and Llama2-13B with RAG (67.7%) on a medical question-answering dataset. Comprehensive evaluations reveal JMLR-13B enhances reasoning quality and reduces hallucinations better than Claude3-Opus. Additionally, JMLR-13B (148 GPU hours) also trains much faster than Meditron-70B (42630 GPU hours). Through this work, we provide a new and efficient knowledge enhancement method for healthcare, demonstrating the potential of integrating retrieval and LLM training for medical question-answering systems.

• The paper presents JMLR, a method that synchronizes LLM and retriever training to enhance reasoning and reduce hallucinations in medical question-answering.
• It employs a dual-parameter optimization with a unique LLM-Rank loss to improve document relevance and answer quality.
• Experimental results across multiple medical datasets show JMLR outperforming larger models with far lower computational costs.

Joint Medical LLM and Retrieval Training for Enhancing Reasoning and Professional Question Answering Capability

This paper introduces JMLR, a novel approach to improving the accuracy and reasoning capabilities of medical LLMs by jointly training LLMs and information retrieval systems. The paper demonstrates JMLR's effectiveness in enhancing medical question-answering systems by reducing the hallucinations common in LLM outputs and optimizing computational efficiency.

Introduction to JMLR

JMLR stands out by integrating the training of retrievers with LLMs, deviating from traditional methods where these components are trained separately. Conventional approaches involve training a retriever to fetch relevant documents and fine-tuning an LLM based on retrieved data (Figure 1). JMLR synchronizes the training of the retriever and the LLM, thus aligning their operations to jointly improve the model's performance in medical question-answering tasks.

Figure 1: Comparison between different domain adaptation methods: traditional domain pretraining method (left), RAG (middle), and JMLR (right). JMLR retrieves the documents to reduce the hallucination.

Implementation Strategy

The paper details the architectural setup for JMLR, emphasizing its dual-parameter optimization process. The optimization strategy involves simultaneously updating both retriever and LLM parameters using a gradient descent approach based on a combined loss function, ensuring the retriever identifies contextually helpful documents that enhance LLM responses.

The implementation leverages the LLM-Rank loss system, a unique mechanism for assessing the impact of retrieved documents on the LLM's answer quality. This allows for the retriever's prioritization of documents based on their utility in improving LLM performance.

Experimental Setup and Results

The authors conducted extensive experiments across multiple datasets, including MMLU-Medical, MedMCQA, MedQA, and Amboss, demonstrating JMLR's superior performance in medical question-answering tasks compared to existing state-of-the-art models like Meditron and ChatGPT.

Figure 2: JMLR achieved the highest average accuracy across the MMLU-Medical, MedMcQA, MedQA, and Amboss datasets, utilizing only 148 GPU hours.

The results indicate that JMLR models, with parameters ranging up to 13 billion, consistently outperform larger models such as Meditron 70B in accuracy and computational efficiency. JMLR reduces training time significantly, requiring only 148 GPU hours compared to Meditron's 42630 GPU hours.

Discussion and Implications

The success of JMLR underscores the importance of integrating document retrieval directly into LLM training processes, especially in domains where accuracy is paramount, like healthcare. By reducing hallucination and enhancing retrieval effectiveness, JMLR offers a computationally efficient solution for deploying medical question-answering systems.

The implications of this approach are vast, potentially transforming how AI supports clinical decision-making by providing reliable, contextually grounded answers, thereby advancing the accessibility of crucial medical knowledge.

Conclusion

JMLR represents a significant advancement in combining retrieval mechanisms with LLM training, showcasing improvements in accuracy, reasoning capabilities, and efficiency. This research opens avenues for further exploration into synchronous training mechanisms in LLMs, with potential applications extending beyond medical domains, relying on robust information retrieval architectures to combat hallucination effects effectively.