Aligning Medical Conversational AI through Online Reinforcement Learning with Information-Theoretic Rewards

TL;DR

This work reframes medical history-taking as an information acquisition problem, enabling a medical conversational AI to learn questioning strategies through online reinforcement learning without relying on expert conversation data. By combining an entropy-based information gain reward with LLM-based quality assessments and optimizing via Group Relative Policy Optimization, the approach guides multi-turn dialogues toward clinically relevant and targeted questioning. Empirical results show IGFT improves HPI extraction accuracy and generalizes from concise Avey vignettes to longer MIMIC histories, outperforming commercial baselines and domain-specific single-turn systems. The findings suggest principled, objective-driven RL can facilitate domain-specific AI alignment in healthcare tasks, though practical deployment requires further validation, efficiency improvements, and careful safety considerations.

Abstract

We present Information Gain Fine-Tuning (IGFT), a novel approach for training medical conversational AI to conduct effective patient interviews and generate comprehensive History of Present Illness (HPI) without requiring pre-collected human conversations. IGFT combines online Group Relative Policy Optimization (GRPO) with information-theoretic rewards, enabling models to learn from self-generated conversations with simulated patients. Unlike existing approaches that rely on expensive expert-annotated conversations or static datasets, our online RL framework allows models to discover effective questioning strategies through exploration. Our key innovation is an information gain reward function that tracks which clinical entities such as symptoms, temporal patterns, and medical history, are revealed during conversation. Each question's reward is computed based on its expected information gain combined with GPT-4o-mini quality assessments across dimensions including clinical relevance, patient engagement, and specificity. This hybrid approach ensures models learn to ask targeted, clinically appropriate questions that efficiently gather diagnostic information. We fine-tune two models using LoRA: Llama-3.1-8B-Instruct and DeepSeek-R1-Distill-Qwen-7B (a reasoning-optimized model). Training exclusively on Avey data containing concise HPIs, we evaluate generalization to MIMIC data with longer, more elaborate HPIs. DeepSeek-R1-Distill-Qwen-7B (IGFT) achieves F1 scores of 0.408 on Avey (10.9% improvement over base) and 0.289 on MIMIC (12.9% improvement), while Llama-3.1-8B-Instruct (IGFT) reaches 0.384 and 0.336 respectively. Both models outperform OpenAI's model on MIMIC and surpass medical domain-specific baselines like HuatuoGPT and UltraMedical, which were optimized for single-turn medical QA rather than multi-turn conversations.

Figures (2)

• Figure 1: Online training framework for medical dialogue. The Doctor LLM generates questions answered by a Patient LLM simulating cases from ground truth HPIs. Information gain rewards are computed by comparing revealed entities against ground truth, combined with LLM-based quality assessment. The model updates through GRPO based on these rewards, creating an online learning loop that improves questioning strategies without requiring human conversation data. The Generated HPI (shown in green) is produced only during evaluation to measure performance, not during training.
• Figure 2: Example of HPI evaluation through atomic statement extraction and matching. Atomic statements are extracted from both ground truth and generated HPIs using GPT-4o-mini. Green indicates correctly captured information, while red indicates missed statements (left) or unnecessary information not present in the reference (right). In this example, the model correctly identified chest pain and dizziness but missed the reduced consciousness level and added denials of fever and cough, which were not part of the ground truth. This yields precision = 2/4 and recall = 2/3.