PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.26587v1cs.CL

TL;DR

The paper addresses the open question of how biomedical KG signals should be integrated into LLMs for clinical diagnosis. It systematically compares KG tasks and training paradigms, then diagnoses their parameter-update geometry. KG-judgment training, especially with KL regularization, produces sparse surgical alignment, while task-specific SFT produces dense updates; sparse-update paradigms also improve reasoning quality.

  • Problem

    How to incorporate biomedical KGs into LLMs for clinical diagnosis remains an open question, despite KGs offering structured scaffolding for grounding reasoning.

  • Method

    The paper systematically studies five KG task formulations, three training paradigms, two KGs, and three base LLMs using transfer evaluations plus GID and GD optimization-geometry diagnostics.

  • Results

    All paradigms improve over the non-finetuned baseline, while KG-judgment training under KL regularization produces sparse localized updates and task-specific SFT produces dense updates.

  • Takeaways & Limitations

    Clinical KG-LLM integration should evaluate optimization geometry alongside task accuracy because sparse updates preserve broader reasoning and can improve reasoning quality despite lower in-domain accuracy.

  • Takeaways & Limitations

    The evaluation covers three LLMs and should be expanded across more diverse domains and additional models; the paper does not propose a new RL framework for improving performance metrics.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies how biomedical KG signals should be integrated into LLMs for clinical diagnosis. It compares training configurations and shows that accuracy alone misses important differences in knowledge transfer and parameter changes.

  • Clinical LLM outputs can contain factual errors, motivating biomedical KGs as structured scaffolding for grounded medical reasoning.
  • The study spans five KG task formulations, three training paradigms, two KGs, and three base LLMs.
  • All paradigms improve over the non-finetuned baseline, but comparable in-domain accuracy can produce substantially different transfer behavior across benchmarks.
  • GID measures update-footprint sparsity and GD measures parameter deviation from the pretrained baseline.
  • Judgment-based KG training with KL regularization yields sparse, localized updates called surgical alignment, whereas task-specific SFT produces dense updates.
  • A controlled ablation finds that the KG objective and KL independently contribute to sparsity, which co-occurs with improved PDSQI-9 reasoning quality even when in-domain accuracy is lower than task-specific SFT.

2 Related Work

Related work integrates KGs either as inference-time context or as training supervision. Recent studies also use reward models as explicit reasoners, while this paper examines KG-LLM integration through training and optimization geometry.

  • KG as Context: Inference-time approaches use KG retrieval, prompting, external memories, candidate ranking, subgraph prompts, or weighted KG construction.
  • KG as Training Supervision: Training-time approaches supervise LLMs with KG structure, parameter-efficient adapters, KG-path rewards, or KG-enhanced fine-tuning.
  • Reward Models as Reasoning: Recent reward-modeling work treats reward models as implicit or explicit reasoning systems rather than mere scorers.
  • KG path extraction: KG path extraction maps note entities to starting concepts and diagnoses to target concepts, labeling paths reaching targets positive and other paths negative.

3 Data Overview

The study uses two biomedical KGs, two clinical training datasets, and two diagnosis-focused medical question-answering datasets. These resources support KG training and downstream clinical evaluation.

  • Knowledge graphs: The study uses UMLS and PrimeKG as its two knowledge graphs.UMLS retains 107 SNOMED CT-based diagnostic relations, while PrimeKG focuses on Disease, Phenotype, and Drug nodes.
  • Training data: ProbSum contains 1,005 physician-annotated ICU progress notes from MIMIC-III.
  • Training data: DDX-Plus is a large-scale synthetic differential-diagnosis dataset covering 49 pathologies, 110 symptoms, and 113 antecedents.
  • Downstream evaluation: The study evaluates diagnosis-related MedQA questions using 1,796 training and 251 test samples.
  • Downstream evaluation: MedMCQA contributes 2,560 diagnosis-related questions from medical entrance examinations.

4 Preprocessing

Preprocessing maps clinical concepts to KG nodes, extracts short reasoning paths, and formulates five KG reasoning tasks. Prompts expose training diagnoses but remove them during evaluation.

  • KG path extraction: Entities in each patient note become starting KG concepts, diagnoses become target concepts, and breadth-first search explores paths up to two hops.
  • KG path extraction: Paths reaching a target concept are labeled positive, while other extracted paths are labeled negative.
  • KG training tasks: The five tasks comprise three path-selection variants—P@10, P@2, and PN@10—and two generative tasks, NHP and PC.
  • KG training tasks: Path Selection judges candidate-path validity, whereas Next-Hop Prediction and Path Completion generate continuations of partial reasoning chains.
  • Prompt structure: Training prompts include system instructions, progress notes, diagnoses, either CANDIDATE_PATH or PARTIAL_PATH, and a task instruction; evaluation removes diagnoses.

5 KG Training Paradigm

The study compares KG-based supervised, multitask, and reward-based training paradigms across multiple task formulations, while preserving the base model distribution with KL regularization. It evaluates lexical, concept-level, exact-match, and diagnostic-reasoning outcomes against task-specific SFT, retrieval, and non-finetuned baselines.

  • Experiments use Qwen2.5-7B-Instruct, Qwen3-8B, and Gemma-7B-IT as base models.
  • KG-based SFT: SFT trains models on each KG task formulation, with multitask SFT pooling all five formulations for more uniform exposure.The multitask data are subsampled to keep per-task training size consistent.
  • KG-based SFT: KL regularization is added to cross-entropy to preserve the non-finetuned model’s output distribution, with λ controlling KL strength.The tested KL-strength range is λ ∈[0.01, 1].
  • KG Reward Model as Reasoning Training: Reward-based training uses GRPO, RM-R1, and Comp-GRPO to optimize preferences over KG-grounded candidate outputs.GRPO compares candidate paths within groups; RM-R1 combines reasoning distillation with GRPO, while Comp-GRPO uses discrete and path-alignment rewards.
  • Evaluation and Baselines: Evaluation combines ROUGE-L, CUI-F, exact-match QA accuracy, and PDSQI-9, using task-specific SFT, RAG, and NFT as baselines.ROUGE-L and CUI-F cover KG-path and diagnosis-prediction tasks, while MedQA and MedMCQA use exact matching.

6 Task-level Results

Task-level evaluations show that KG training improves over the non-finetuned baseline, but transfer depends strongly on task formulation and training paradigm. Multi-task and KG-trained models show uneven generalization across graph tasks, diagnosis prediction, medical QA, and clinical reasoning quality.

  • Overall task-level performance: All KG training paradigms improve over the non-finetuned baseline, while comparable in-domain accuracy can conceal different transfer behavior.Table 2 compares KG-based SFT, GRPO, and RM-R1 against the NFT baseline.
  • Path judging vs. generation: Path-judging tasks generalize well within their family but transfer poorly to generative tasks such as NHP and PC, and vice versa.This asymmetry indicates limited cross-family generalizability of path discrimination objectives.
  • Multi-task robustness and training paradigm: Multi-task SFT provides the most consistent cross-task performance across both UMLS and PrimeKG, improving robustness mainly within judgment-based formulations.Its generalization to generative tasks remains limited.
  • Downstream task performance: Task-specific SFT achieves the highest in-domain diagnosis-prediction performance on ProbSum and DDXPlus, while reward-trained models outperform NFT but not task-specific SFT.Diagnosis prediction is evaluated with ROUGE-L and CUI-F on the two test sets.
  • Downstream task performance: Medical-QA results are mixed: task-specific SFT wins on MedMCQA, whereas KG-trained models match or exceed it on MedQA.Table 4 notes that multiple methods can achieve the highest metric because performance is similar across methods.
  • Clinical reasoning quality: KG-trained models improve organization, comprehensibility, and synthesis on PDSQI-9 for Qwen7B and Gem7B, without consistent gains in extractive accuracy or thoroughness.The gains are structural rather than surface-level.

7 Gradient Analysis

The analysis measures layer-wise update direction, magnitude, and discrete intervention scope relative to the non-finetuned baseline. KG-judgment objectives yield sparse, localized updates associated with stronger reasoning improvements, while task-specific SFT produces broad, distorted updates; the ablation attributes sparsity independently to the objective and KL regularization.

  • 7.1 Diagnostic Framework: GID measures the fraction of layer components touched by updates exceeding δ = 10^-9, while GD measures directional misalignment as 1 − Cos.The intervention threshold separates optimization updates from floating-point noise, and GID captures scope rather than update magnitude.
  • 7.2 Results: Task-specific SFT modifies nearly every component at every layer with GD near 1.0, whereas KG-guided methods show lower distortion and more conservative updates.RMR1 falls between these regimes, consistent with its lack of explicit KL regularization.
  • 7.2 Results: The sparse-update ordering holds across Gemma-7B, Qwen2.5-7B, and Qwen3-8B, motivating the term surgical alignment for localized changes that leave most pretrained parameters intact.The pattern is reported as reflecting the training paradigm rather than the base model.
  • 7.2 Results: The sparsest GID profiles, produced by KG-SFT and GRPO, coincide with improvements on higher-order PDSQI-9 dimensions, unlike task SFT despite its in-domain ROUGE-L advantage.The reported dimensions include organization, synthesis, and related higher-order reasoning qualities.
  • 7.3 Ablation: objective vs. KL divergence: At λ = 0, ProbSum SFT reaches 22.92 ROUGE-L versus 18.88 for KG SFT, while adding KL reduces ProbSum SFT performance across all reported evaluations.ProbSum ROUGE-L changes from 22.92 to 19.48, MedQA from 64.54 to 63.61, and MedMCQA from 59.18 to 57.82.
  • 7.3 Ablation: objective vs. KL divergence: KG SFT remains essentially flat across λ, because its objective already creates sparse updates without KL; KL primarily constrains the broad footprint of ProbSum SFT.KG SFT is reported as nearly unchanged across ProbSum, MedQA, and MedMCQA evaluations, whereas KL reduces task-SFT in-domain accuracy.

8 Conclusion

The paper identifies surgical alignment through GID and GD as sparse, localized updates from judgment-based KG training that preserve broader pretrained reasoning. It concludes that optimization geometry should complement task accuracy when evaluating clinical KG–LLM integration.

  • 8 Conclusion: GID and GD identify surgical alignment as sparse, localized parameter updating associated with preservation of the pretrained model’s broader reasoning.The conclusion presents this as a property of judgment-based KG training.
  • 8 Conclusion: Optimization geometry should be evaluated alongside task accuracy for clinical KG–LLM integration.

10 Limitations

The evaluation spans multiple task formulations, training paradigms, and KGs but remains limited in domain and model coverage. The paper also does not introduce a novel RL framework for improving performance metrics.

  • 10 Limitations: The evaluation is limited to three commonly used LLMs and could be expanded across additional models and diverse domains.
  • 10 Limitations: The study introduces a mechanistic GID perspective but does not propose a novel RL training framework aimed at improving performance metrics.

11 Ethical Considerations

The study uses deidentified MIMIC-III-derived data in a HIPAA-compliant Azure environment and evaluates reasoning with LLM-based judges. Its KG concept metric compares predicted and gold concept sets through precision, recall, and CUI-F.

  • 11 Ethical Considerations: ProbSum is built from deidentified MIMIC-III data containing no personally identifiable information, with Azure described as HIPAA-compliant.
  • 11 Ethical Considerations: Azure GPT-o3-mini generates reasoning traces, while Azure GPT-5-mini scores PDSQI-9 evaluations with ICC ≥ 0.8.
  • 11 Ethical Considerations: CUI-F evaluates concept grounding by comparing predicted KG concepts with gold concepts through precision and recall.

C Detailed training dynamics

The gradient analysis compares fine-tuned models with non-fine-tuned baselines using four metrics and reveals distinct optimization behaviors across training paradigms. Task SFT makes aggressive changes, whereas multi-task methods preserve alignment and adapt through smaller updates.

  • Gradient metrics: Cosine Similarity, Mean Squared Error, Energy Shift, and Sign Agreement characterize the optimization trajectory relative to the non-fine-tuned baseline.Cos measures gradient alignment, MSE measures update differences, Egy captures logarithmic magnitude change, and Sgn measures retained parameter polarity.
  • Task SFT: Cos ≈0.11–0.42 and significantly higher Mse mark Task SFT as an aggressive parameter-modification strategy.The analysis associates this behavior with potentially overwriting pretrained features needed for general capability.
  • Multi-task learning: Cos > 0.96 and Sgn > 0.94 show that MT SFT and MT GRPO remain closely aligned with the pretrained model.Their small Egy values indicate precise, low-rank adjustments for integrating KG knowledge rather than wholesale retraining.
  • Reward modeling: MT RM-R1 exhibits moderate structural deviation without an explicit KL divergence penalty, while retaining much of the model’s general knowledge.The passage characterizes reward modeling as reshaping the feature space without substantially compromising general knowledge.
  • Implications: The KG-guided methods update parameters in a small task-relevant subspace, mitigating catastrophic forgetting while supporting downstream adaptation.Figures 6 and 7 and Table 7 are presented as evidence of distinct optimization behaviors across the training paradigms.
Loading 2608.26587v1…