Surgical Alignment in Knowledge Graph Training for Clinical Diagnosis with Large Language Models

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

cs.CLAug 27, 2026

Surgical Alignment in Knowledge Graph Training for Clinical Diagnosis with Large Language Models

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

Organizations: University of Colorado Anschutz Medical Campus · University of Wisconsin Madison · Loyola University

Abstract

Biomedical knowledge graphs (KGs) offer structured medical knowledge that can ground large language model (LLM) reasoning in clinical diagnosis application, yet how KG signal should be integrated into LLMs remains an open question. We present a systematic study spanning five KG task formulations, three training paradigms, two KGs, and three base LLMs. At the task level, all paradigms improve over the non-finetuned baseline, but methods with comparable in-domain accuracy show substantially different knowledge transfer behavior. We introduce Gradient Intervention Density (GID) and Gradient Distortion (GD) to measure how broadly an optimizer modifies the pretrained model. GID and GD together reveal a clear divide: KG-judgment training under KL regularization produces sparse, localized updates (a regime we term as surgical alignment), while task-specific SFT produces dense ones. A controlled ablation shows that the objective and KL contribute to sparsity independently, and the paradigms that produce sparse updates also improve reasoning quality, even when their in-domain accuracy is lower than task-specific SFT. Assessing KG-LLM integration thus requires complementing accuracy with optimization-geometry diagnostics. Our implementation can be found at https://github.com/LARK-NLP-Lab/Surgical-Alignment.

Figures & tables

Figure 1: KG path extraction from patient progress notes. The concept extractor ( QuickUMLS Soldaini (2016) for UMLS KG and Simstring-fast Okazaki and Tsujii (2010) for PrimeKG) identifies the starting concepts based on entities in progress notes, and gold concepts based on diagnoses labels. For example in this figure, we can see that starting from concept CUI-2 , the paths which end with CUI-X and CUI-Y are positive samples and other paths are negative.

Figure 2: Training Prompt structure for the different task formulation. The basic format includes the system initialization explaining the scenario, patient progress notes, patient diagnoses, either CANDIDATE_PATH or PARTIAL_PATH based on the task, and the final line explaining the task. For evaluation, we follow the same structure of the prompt, but remove the patient diagnoses part. CUI denotes Concept Unique Identifier, representing the KG node which our model’s output maps to.

Task Type	Task	Description
Path Selection	P@10	Given 10 candidate KG paths, identify one valid path.
	P@2	Given 2 candidate KG paths, identify one valid path.
	PN@10	Given 10 candidate KG paths, identify multiple valid paths.
Next-Hop	NHP	Given a partial KG path, predict the next hop.
Path Completion	PC	Given a partial KG path, predict the remaining path.

Table 1: Task formulations for training LLMs to reason over medical knowledge graphs. The three task types (Path Selection, NHP, and PC) probe complementary reasoning abilities. The Path Selection task requires model to judge the validity of the reasoning paths, whereas NHP and PC are generative reasoning tasks which require the models to complete the partial reasoning chains.

Figure 3: Cross-task generalization among the path-judging formulations across UMLS and PrimeKG. Results are reported in ROUGE-L scores. Due to space constraints, we present results from Qwen2.5-7B-Instruct, but observe consistent patterns across the other LLMs.

	UMLS						PrimeKG
Training Paradigm	Qwn7B		Qwn8B		Gem7B		Qwn7B		Qwn8B		Gem7B
Training Paradigm	RL	CF	RL	CF	RL	CF	RL	CF	RL	CF	RL	CF
NFT	46.14	51.29	36.42	39.26	24.70	28.84	49.55	37.23	32.70	25.03	45.99	40.37
SFT	59.89	60.92	44.71	45.90	68.03	69.68	63.34	64.73	47.16	41.90	71.95	65.23
GRPO	57.74	59.24	56.82	58.29	64.37	66.34	68.09	60.95	49.66	40.81	71.73	65.75
Comp-GRPO	59.30	61.60	53.40	54.62	61.23	63.08	68.62	61.13	51.93	43.95	71.89	65.75

Table 2: Rouge-L and CUI-F scores comparing various graph training paradigm across Qwen2.5-7B-Instruct, Qwen3-8B, and Gemma-7B-IT. We compare the multi-task checkpoints for the KG-based SFT, GRPO, and RM-R1 training paradigms against the NFT baseline. GRPO is performed to models which have already undergone SFT on KG paths. Comparison between different task formulations is highlighted in figure 3 . The best performing checkpoints for each KG are marked with bold text.

	Training	ProbSum						DDXPlus
		Qwn7B		Qwn8B		Gem7B		Qwn7B		Qwn8B		Gem7B
		RL	CF	RL	CF	RL	CF	RL	CF	RL	CF	RL	CF
Baseline	NFT	22.14	26.05	04.15	16.36	20.20	25.33	10.90	09.50	01.58	05.72	06.38	05.52
	SFT	26.09	27.53	04.72	17.87	17.48	30.42	22.98	46.20	13.62	38.95	28.60	46.07
	Cross SFT	03.89	08.23	03.79	10.43	02.91	03.76	12.52	08.98	01.62	05.95	04.60	04.72
	RAG - UMLS	14.09	06.09	09.08	03.05	08.49	14.65	05.45	02.66	04.00	01.59	03.37	07.17

Table 3: Rouge-L (RL) and CUI-F (CF) scores for diagnosis prediction on ProbSum and DDXPlus test set. Multi-task checkpoints are evaluated alongside baseline methods. Best performing KG-trained method for each KG is marked as bold.

Task	Training	Qwn7B	Qwn8B	GemB
MedMCQA	NFT	57.14	36.73	36.73
	RAG - UMLS	54.42	36.73	36.73
	RAG - PrimeKG	54.42	36.73	36.73
	Probsum SFT	63.27	36.73	32.65
	DDXPlus SFT	63.27	32.65	30.61
	Best UMLS model	60.20	36.73	36.73

Table 4: QA results on MedMCQA and MedQA dataset. Best performing training paradigm for each model from Table 3 are evaluated alongside NFT models. The highest metric for each model is marked as bold, however the performance across different methods are similar and multiple methods achieve highest metric.

PDSQI-9	NFT			Task SFT			Best UMLS model from Table 3			Best PrimeKG model from Table 3
PDSQI-9	Qwn7B	Qwn8B	Gem7B	Qwn7B	Qwn8B	Gem7B	Qwn7B	Qwn8B	Gem7B	Qwn7B	Qwn8B	Gem7B
Acc.	1.69	2.22	1.52	1.81 $\uparrow$	1.52	2.06 $\uparrow$	1.62	2.12	1.68 $\uparrow$	1.88 $\uparrow$	1.80	1.32
Thorou.	2.14	2.08	1.54	1.55	1.32	1.38	2.16 $\uparrow$	1.80	1.56 $\uparrow$	2.06	1.84	1.44
Useful.	2.46	2.48	1.72	1.93	1.42	1.48	2.38	2.12	1.78 $\uparrow$	2.46	2.2	1.72
Org.	4.12	2.70	2.94	1.95	1.22	1.42	3.64	2.50	3.00 $\uparrow$	4.02	2.16	3.40 $\uparrow$
Comp.	4.50	3.70	4.52	3.28	1.98	2.58	4.32	3.56	4.62 $\uparrow$	4.46	3.14	4.48

Table 5: Best KG-trained LLM on ProbSum, evaluated by PDSQI-9 criteria, consisting of Accuracy (Extractive), Thoroughness, Usefulness, Organization, Comprehensibility, Succinctness, Synthesis/Abstraction metrics. Baselines include NFT models and Probsum SFT models. We compare the baselines with the best performing KG-trained models on the Probsum task, on Table 3 .

\uparrow

denotes an improvement over the baseline models. In case of UMLS KG, the best checkpoints for Qwn7B, Qwn8B, and Gem7B are GRPO, RM-R1, and SFT respectively. In case of PrimeKG, the best checkpoints for Qwn7B, Qwn8B, and Gem7B are GRPO, RM-R1, and RM-R1 respectively.

Figure 4: Gradient analysis across Models and training paradigms. We map the layer-wise gradients using Layer Depth (X), Gradient Distortion (Y), and Gradient Intervention Density (Z).

		KG SFT		ProbSum SFT
Task	$\lambda$	Rouge-L	CUI-F	Rouge-L	CUI-F
Probsum	0.0	18.88 $\pm$ 0.03	24.00 $\pm$ 0.00	22.92 $\pm$ 0.01	25.87 $\pm$ 0.00
Probsum	0.1	18.94 $\pm$ 0.02	22.66 $\pm$ 0.00	19.48 $\pm$ 0.03	24.88 $\pm$ 0.00
		Out-of-distribution test set
		Accuracy		Accuracy
MedQA	0.0	64.14 $\pm$ 1.17		64.54 $\pm$ 0.65

Table 6: Ablation study on Qwen2.5-7B-Instruct model. We report the average ROUGE-L score

\pm

standard deviation (sd) across three runs on the Probsum diagnosis task. Furthermore, we also report mean accuracy scores

\pm

sd across three runs on the MedQA and MedMCQA dataset.

Figure 5: Gradient analysis across Qwn7B models in our ablation study.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: Gradient-level optimization dynamics. Layer-wise analysis for Baseline and UMLS-based methods across three models.

Panel of Figure 7 — Figure 7: Gradient-level optimization dynamics. Layer-wise analysis for Baseline and PrimeKG-based methods across three models.

KG	Method	Met.	Qwen2.5-7B				Qwen3-8B				Gemma-7B
KG	Method	Met.	Q	K	V	O	Q	K	V	O	Q	K	V	O
Baseline	Task SFT	Cos	0.4274	0.4253	0.5275	0.4544	0.3784	0.3424	0.4366	0.3597	0.1137	0.1224	0.1929	0.1470
		Mse	8.9e-08	6.3e-07	5.1e-07	8.5e-08	7.4e-08	3.1e-07	2.7e-07	7.6e-08	1.4e-07	1.4e-07	1.3e-07	1.4e-07
		Egy	-5.2e-09	1.9e-09	1.2e-09	8.5e-10	-2.5e-10	2.3e-09	-1.3e-09	-8.3e-11	-1.8e-09	-2.1e-09	-1.6e-09	-4.3e-10
		Sgn	0.6262	0.6282	0.6443	0.6400	0.6022	0.5990	0.6126	0.6177	0.5354	0.5349	0.5462	0.5479
UMLS	MT SFT	Cos	0.9845	0.9821	0.9886	0.9901	0.9819	0.9782	0.9857	0.9875	0.9881	0.9883	0.9920	0.9920

Appendix

Table 7: Layer-averaged divergence metrics for attention components (Q, K, V, O) across three models (Qwen2.5-7B, Qwen3-8B, Gemma-7B). We compare the baseline Task SFT against three multi-task paradigms (MT SFT, MT GRPO, MT RM-R1) on two KGs (UMLS, PrimeKG).

KG	Method	Threshold	Qwen2.5-7B				Qwen3-8B				Gemma-7B
KG	Method	Threshold	Q	K	V	O	Q	K	V	O	Q	K	V	O
Baseline	Task SFT	$10^{-6}$	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00
		$10^{-7}$	0.04	0.06	0.05	0.06	0.06	0.05	0.04	0.06	0.05	0.05	0.02	0.04
		$10^{-8}$	0.40	0.46	0.44	0.43	0.38	0.43	0.46	0.42	0.41	0.40	0.39	0.36
		$10^{-9}$	0.40	0.46	0.44	0.43	0.38	0.43	0.46	0.42	0.41	0.40	0.39	0.36
		$10^{-10}$	0.40	0.46	0.44	0.43	0.38	0.43	0.46	0.42	0.41	0.40	0.39	0.36

Appendix

Table 8: Gradient intervention density across various thresholds.

LLM Reasoning with Graphs

Large language model reasoning with graphs uses structured relationships to support multi-step inference, evidence retrieval, and planning.

See all 211 papers

Knowledge-Graph Grounding Helps LLMs Only for Out-of-Training Knowledge: A Controlled Study on Clinical Question Answering
Jun 21, 2026Madhulatha Mandarapu, Sandeep KunkunuruFactual Knowledge in Language Models LLM Reasoning with Graphs
samyama
The Graph Language: How Knowledge Graphs Speak to Large Language Models
Aug 2, 2026Giuseppe PirròLLM Grounding Knowledge Augmentation for Language Models
DeMaCS, University of Calabria
MedThink: Enhancing Diagnostic Accuracy in Small Models via Teacher-Guided Reasoning Correction
Apr 9, 2026Xinchun Su, Chunxu Luo, Lipeng Ma +2Language Model Distillation Knowledge Distillation
School of Computer Science, Fudan University, Shanghai, China

Jun 21, 2026J/K move · Enter open · S save

Knowledge-Graph Grounding Helps LLMs Only for Out-of-Training Knowledge: A Controlled Study on Clinical Question Answering

Madhulatha Mandarapu, Sandeep Kunkunuru

samyama

Factual Knowledge in Language Models LLM Reasoning with Graphs Clinical QA LLM Grounding Clinical Language Model Evaluation Biomedical KGs+1

Open paper AskPDF arXiv

A recent Nature Medicine study reports that general-purpose frontier LLMs outperform specialized retrieval-augmented clinical tools on medical benchmarks, and that retrieval can hurt strong models. We ask the natural follow-up: does structured knowledge-graph (KG) grounding change this, and when does grounding help at all? We contribute two results. First, a reproduction: the study's headline HealthBench score (~88) is the Consensus variant, not full HealthBench, where frontier models and ideal completions both score ~46-47 under a physician-calibrated grader (agreement 82.5%); we reproduce GPT-5.2 Consensus =90.9 and flag a score-deflating grader bug. Second, a knowledge-boundary result. Using a graph+vector engine (samyama-graph) over the public biomedical KG PrimeKG, neither naive triple retrieval nor an agentic natural-language-to-Cypher loop (82% successful queries) improves MedQA across a weak-to-strong model ladder (all |Delta| <= 3.4). On a synthetic counterfactual KG, and on a hybrid benchmark mixing known and novel facts, the identical pipeline lifts out-of-training accuracy from chance to ~100% (+68 to +79) while adding nothing on known facts (a no-LLM arm answers both). Across three regimes (no-knowledge, graph-aided, hybrid), grounding helps only insofar as the decisive fact lies outside the model's training -- public-KG facts are redundant, private and novel data are where it pays -- matching the study's institutional-data caveat.

Abstract

LLM Reasoning with Graphs

Explore similar work

Knowledge-Graph Grounding Helps LLMs Only for Out-of-Training Knowledge: A Controlled Study on Clinical Question Answering

The Graph Language: How Knowledge Graphs Speak to Large Language Models

MedThink: Enhancing Diagnostic Accuracy in Small Models via Teacher-Guided Reasoning Correction