Brittleness and Promise: Knowledge Graph–Based Reward Modeling for Diagnostic Reasoning

Saksham Khatwani, He Cheng, Majid Afshar, Dmitriy Dligach, Yanjun Gao

Keywords: Large Language Models, Knowledge Graph, Diagnostic Reasoning, Reward Modeling, Clinical NLP

Abstract:

Large language models (LLMs) show promise for diagnostic reasoning but often lack reliable, knowledge-grounded inference. Knowledge graphs (KGs), such as the Unified Medical Language System (UMLS), offer structured biomedical knowledge that can support trustworthy reasoning. Prior approaches typically integrate KGs via retrieval-augmented generation or fine-tuning, inserting KG content into prompts rather than enabling structured reasoning. We explore an alternative paradigm: treating the LLM as a $\textit{reward}$ model of KG reasoning paths, where the model learns to $\textit{judge}$ whether a candidate path leads to correct diagnosis for a given patient input. This approach is inspired by recent work that leverages reward training to enhance model reasoning abilities, and grounded in computational theory, which suggests that $\textit{verifying}$ a solution is often easier than generating one from scratch. It also parallels physicians’ diagnostic assessment, where they $\textit{judge}$ which sequences of findings and intermediate conditions most plausibly support a diagnosis. We first systematically evaluate five task formulation for knowledge path judging and eight training paradigm. Second, we test whether the path judging abilities generalize to downstream diagnostic tasks, including diagnosis summarization and medical question answering. Experiments with three open-source instruct-tuned LLMs reveal both promise and brittleness: while specific reward optimization and distillation lead to strong path-judging performance, the transferability to downstream tasks remain weak. Our finding provides the first systematic assessment of "reward model style" reasoning over clinical KGs, offering insights into how structured, reward-based supervision influences diagnostic reasoning in GenAI systems for healthcare. Our implementation can be found at
https://github.com/LARK-NLP-Lab/kg-rl-reasoner.

Submission Number: 72