[Paper Review] JMLR: Joint Medical LLM and Retrieval Training for Enhancing Reasoning and Professional Question Answering Capability

Junda Wang, Zhichao Yang+2 more|arXiv (Cornell University)|Feb 27, 2024

Artificial Intelligence in Healthcare and Education11 citations

TL;DR

The paper introduces Joint Medical LLM and Retrieval Training (JMLR), a synchronized retriever-LLM training framework that improves medical QA and reasoning while reducing training time. It reports state-of-the-art open-source results on medical benchmarks with 7B and 13B Llama-based models.

ABSTRACT

Large Language Models (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.

Motivation & Objective

• Motivation to improve medical QA and reasoning with domain-specific knowledge.
• Address hallucinations by grounding LLMs with retrieved medical guidelines and texts.
• Develop a joint training paradigm that updates retriever and LLM together for better alignment.
• Evaluate efficiency gains and performance against traditional pretraining+finetuning pipelines.

Proposed method

• Use Shifted Sparse Attention (S2-Attn) to handle long input contexts.
• Employ ColBERT-based retriever with a joint LLM-retriever training objective (LLM-Rank loss).
• Train on QA pairs from AMBOSS and USMLE with top retrieved documents fed into the LLM.
• Compute LLM-driven loss and update retriever parameters via a rank-based signal that reflects LLM improvement.
• Sample top-30 retrieved docs per iteration and feed top-7 into the LLM for answer generation and reasoning.
• Compare integrated JMLR with separate RAG and with baseline pretraining/fine-tuning across 7B and 13B Llama models.

Experimental results

Research questions

• RQ1 Does synchronizing retriever and LLM training improve medical QA accuracy and reasoning compared with conventional pretraining-finetuning and RAG baselines?
• RQ2 Can JMLR reduce training time and resource use while maintaining or surpassing state-of-the-art performance on medical benchmarks?
• RQ3 How does JMLR affect hallucination propensity and explainability in medical QA?
• RQ4 What is the impact of model scale (7B vs 13B) on JMLR's performance and efficiency?

Key findings

• JMLR-13B achieves 81.2% on AMBOSS and 61.3% on MedQA, outperforming Meditron-70B and ChatGPT on these datasets.
• JMLR-7B achieves 68.7% on AMBOSS and 51.7% on MedQA, surpassing several public baselines.
• Training time with JMLR is 37 hours, significantly shorter than the conventional pretraining (127h) plus finetuning (17h).
• JMLR-7B and JMLR-13B show strong results across MMLU-Medical, MedMCQA, MedQA, and AMBOSS benchmarks, indicating improved medical reasoning and QA capability.
• GPT-4 and three physicians independently judged JMLR-13B’s reasoning as superior in most cases (GPT-4 win rate 0.63; experts 0.60).