RareDx: Controlled Knowledge Integration and Graph-Grounded Policy Optimization for Rare-Disease Diagnosis

Bo Zhang; Yuchen Wang; Dongbai Li; Matthew Yu Heng Wong; Qingkai Zeng; Lijun Wang; Tien-Yin Wong; Peng Cui

cs.AISep 28, 2026

RareDx: Controlled Knowledge Integration and Graph-Grounded Policy Optimization for Rare-Disease Diagnosis

Authors: Bo Zhang, Yuchen Wang, Dongbai Li, Matthew Yu Heng Wong, Qingkai Zeng, Lijun Wang, Tien-Yin Wong, Peng Cui, +1 more

Organizations: Xi’an Jiaotong University · UIUC · Tsinghua University · University of Cambridge · Nankai University · Zhejiang University · Yale University

Abstract

Rare-disease diagnosis is a long-tail reasoning problem: phenotypes are incomplete, individual disorders are sparsely documented, and relevant evidence is distributed across ontologies, gene annotations, and biomedical text. Language models consequently favor common conditions, miss rare candidates, or produce plausible but invalid names. We introduce RareDx, which couples controlled evidence use with knowledge-graph-grounded policy optimization. RareDx-Harness normalizes heterogeneous records into one ranked-diagnosis task and compares direct inference, static retrieval, adaptive tools, and structured phenotype-gene-disease reasoning over a shared knowledge layer. The training pipeline combines Top-10 post-training with RareDx-KGPO, our knowledge-graph-grounded policy optimization method. Its reward projects predictions into a canonical disease graph and integrates curated graded relevance, ontology proximity, biomedical similarity, and phenotype consistency. Vocabulary and output-budget constraints prevent dense partial credit from rewarding fabricated or overlong differentials. Across eight benchmarks, the complete RareDx system centered on Qwen3.5-9B reaches 38.34 macro Hit@10, 1.60 points above GPT-5.5 under the archived protocol; a disjoint validation-selection audit retains a 6.80-point routing gain over Direct on held-out cases. The 27B system reaches 23.53/36.56/40.76 at Hit@1/5/10. Controlled ablations show that retrieval is not uniformly helpful and that controlled routing is central to the gain. These results indicate that structured medical knowledge can turn a compact model into a competitive diagnostic ranker across heterogeneous long-tail settings in clinical practice.

Figures & tables

Figure 1: RareDx-Harness. Patient cases from heterogeneous benchmarks are normalized into a common ranked-diagnosis task. The controlled strategy matrix compares direct inference, one-shot static retrieval, adaptive ReAct tool use, and structured phenotype-to-gene-to-disease reasoning over the same local medical knowledge infrastructure. Canonicalization and rank aggregation are shared across strategies to preserve controlled comparison across all branches.

Figure 2: Two-stage post-training. Top-10 SFT establishes the output contract and broad diagnostic knowledge. RareDx-KGPO then optimizes ranked differential diagnoses using disease-graph, semantic, and phenotype evidence, with explicit safeguards against invalid names and reward hacking.

(a) Hit@1: first-ranked diagnosis
Model	MyGene2	RAMEDIS	MME	HMS	LIRICAL	RDS	RDC	Pheno.	Macro
Closed-source LLMs
Claude Opus 4.7	9.6	21.1	12.5	22.2	23.4	12.4	15.2	25.6	17.8
GLM-5.2	11.6	18.5	28.3	21.3	26.1	11.6	15.5	22.8	19.5
GPT-5.5	20.5	21.5	29.5	23.8	27.2	10.7	12.1	30.8	22.0
Open-weight baselines

Table 1: Complete diagnostic ranking results (%) on eight benchmarks. The vertically stacked panels report Hit@1, Hit@5, and Hit@10 within one unified table. Macro is the unweighted mean across datasets. All outputs share one disease normalizer; bold marks panel-best values.

Figure 3: Reward dynamics for rare-disease post-training from a shared initialization. Vanilla GRPO improves initially but collapses after approximately 210 steps, consistent with reward hacking through indiscriminate disease-name generation. DAPO and OPD avoid catastrophic collapse but converge to lower rewards. RareDx-KGPO (Ours) supports sustained improvement and remains stable at the final checkpoint without late-stage collapse.

Figure 4: Per-dataset effect of applying the archived exploratory RareDx-Harness route to the same knowledge-grounded checkpoint. Cells show the absolute change in Hit@1, Hit@5, and Hit@10. Gains concentrate on RAMEDIS and HMS at both scales, with an additional large 27B gain on RDS. Near-zero cells correspond to retaining Direct when the audit provided no measured benefit; the small LIRICAL regression illustrates that routing is not error-free.

Panel	Configuration	Hit@1	Hit@5	Hit@10
A	Direct	17.03	26.63	31.27
A	Adaptive ReAct	16.49	27.51	33.53
B	Direct	17.2	23.7	27.5
B	3-hop, dense gene retrieval	16.2	23.6	28.4
B	3-hop, HPO-Resnik	22.2	31.7	36.3

Table 2: Component ablations with fixed backbone and judge. Panel A uses Qwen3.6-Flash and macro averages over seven benchmarks. Panel B uses Qwen3.6-35B-A3B on Phenopackets.

Figure 5: Observed accuracy by retrieval depth for 6,230 ReAct trajectories, with 95% Wilson intervals Wilson (1942) and stratum sizes for binomial testing. Depth is policy-selected, so the comparison describes behavior rather than estimating the causal effect of retrieval depth.

System	Hit@1	Hit@5	Hit@10
Qwen3.8-27B, direct	2	12	28
Ours 9B, direct	4	12	16
Ours 9B, full Harness	14	38	46

Table 3: Paired audit on the first 50 RAMEDIS records. Values are percentages under the disease matcher and parsing contract used throughout.

Figure 6: Paired rank-category transitions on the fixed 50-case RAMEDIS audit. Each ribbon follows one patient from direct Qwen3.8-27B to RareDx-9B with the full Harness. RareDx recovers 16 diagnoses missed by the larger model, 13 of which enter the Top-5.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Cases	Dataset	Cases
MyGene2	146	RAMEDIS	624
MME	40	HMS	88
LIRICAL	370	RDS	1,803
RDC	678	Phenopackets	500
Total			4,249

Appendix

Table 4: Number of test records contributing to the macro average. RareBench coverage includes all questions, reported separately for RAMEDIS, MME, HMS, and LIRICAL.

Model	Dataset	$n_{dev}$	$n_{test}$	D, dev	A, dev	Pick	D, test	A, test
9B	MyGene2	33	113	24.2/27.3/39.4	12.1/18.2/27.3	D	21.2/29.2/41.6	14.2/29.2/32.7
9B	RAMEDIS	113	511	11.5/18.6/21.2	16.8/40.7/47.8	A	9.6/14.9/18.4	17.6/37.0/41.1
9B	MME	9	31	0.0/11.1/11.1	0.0/0.0/0.0	D	6.5/9.7/9.7	12.9/12.9/12.9
9B	HMS	22	66	22.7/27.3/27.3	27.3/45.5/59.1	A	9.1/12.1/13.6	24.2/37.9/47.0
9B	LIRICAL	62	308	17.7/27.4/29.0	16.1/27.4/30.7	A	22.1/30.8/32.5	19.2/27.0/30.8
9B	RDS	342	1,461	11.7/27.2/38.3	–	D	10.1/24.2/37.2	–

Appendix

Table 5: Validation-only route audit. Each metric cell is Hit@1/5/10 (%). D and A denote Direct and Diagnostic Audit. A dash means that a complete archived Audit candidate was not available, so the candidate set contained Direct only.

Stage	Component	Rows	Content
SFT	train	51,368	RareArena cases with constructed Top-10 targets
SFT	validation	932	held-out Phenopacket-format prompts
RL	RareArena	42,977	article-backed case reports and test results
RL	Phenopacket	8,392	phenotype/gene disease prompts
RL	gene-phenotype hard set	3,928	synthetic hard cases
RL	simulated HPO subset	2,714	synthetic incomplete-phenotype cases

Appendix

Table 6: Training and validation artifacts used by the reported 9B pipeline.

Reported row	Request alias	Interface	Decoding and tools
GPT-5.5	gpt-5.5	OpenAI-compatible proxy, xiaoai.plus/v1	temperature, top- $p$ , seed, and output cap: provider default; tools: none
Claude Opus 4.7	claude-opus-4-7	OpenAI-compatible proxy, xiaoai.plus/v1	temperature, top- $p$ , seed, and output cap: provider default; tools: none
GLM-5.2	glm-5.2	Alibaba MaaS compatible API for phenotype sets; DashScope Generation for RDS/RDC	decoding and output cap: provider default; tools: none

Appendix

Table 7: Recorded proprietary-baseline invocation. “Provider default” means the argument is absent from the archived request rather than assigned an inferred value.

Method	Update	Signal	Teacher
GRPO	Group relative	Task reward	No
DAPO	Dynamic groups, asymmetric clip	Task reward	No
OPD	On-policy distillation	Token distributions	Yes
RareDx-KGPO (Ours)	Group relative	Medical graph reward	No

Appendix

Table 8: Optimization and supervision differences in the RL comparison.

Model	Observed output
Qwen3.8-27B	Restates phenotype categories and surveys metabolic mechanisms, but does not produce a disease-ranked answer before truncation. MCADD is not recovered.
Ours 9B + Harness	Ranks MCADD first, followed by carnitine palmitoyltransferase II deficiency, very-long-chain acyl-CoA dehydrogenase deficiency, multiple acyl-CoA dehydrogenase deficiency, and carnitine-acylcarnitine translocase deficiency.

Appendix

Table 9: Observed outputs for RAMEDIS case 12.

Direct model	Complete block	Mean words	Median words
Qwen3.8-27B	0/50	462.2	451.0
Ours 9B	50/50	91.3	90.5

Appendix

Table 10: Direct-generation behavior on the 50-case audit. Answer-block compliance requires both opening and closing tags. Output words are counted after whitespace tokenization and are reported only as a model-agnostic verbosity diagnostic.

Figure 7: Controlled reward audit on 120 held-out diagnoses. Panel A removes the disease-phenotype graph signal while holding outputs fixed. Panel B removes the vocabulary gate for pseudo-diseases or the Top-10 budget for a delayed exact match.

Probe	Full	Ablated	Positive, full	Positive, ablated
Exact diagnosis	1.000	1.000	100.0	100.0
Phenotype neighbor	0.316	0.008	100.0	0.8
Pseudo-disease	-0.175	0.320	0.0	82.5
Gold at rank 11	0.068	0.279	100.0	100.0
Malformed output	0.000	0.000	0.0	0.0

Appendix

Table 11: Numerical reward audit. “Ablated” removes the phenotype-graph channel, vocabulary gate, or Top-10 budget for the corresponding probe. Positive is the percentage of constructed outputs receiving reward greater than zero.

Figure 8: List-quality audit on 120 held-out diagnoses. Every list contains the exact diagnosis at rank 1. Panel A sweeps the nDCG weight with case-level bootstrap 95% intervals. Panel B reports within-case reward variation across coherent, reversed, and unrelated differentials.

Calls	Cases	Hit@1	Hit@5	Hit@10
0	324	2.16	3.40	4.01
1	383	35.77	54.31	57.96
2	1,423	38.65	54.25	59.31
3	4,100	22.10	37.90	43.32

Appendix

Table 12: ReAct diagnostic performance stratified by the number of retrieval calls selected by the policy. Values are percentages.

Strategy	LLM generations	Local retrieval	External API
Direct	1	None	No
Static RAG	1	Once	No
Adaptive ReAct	$T$	Adaptive	No
Structured 3-hop	1	HPO + gene lookup	No
RRF fusion	$M$	Strategy-dependent	No

Appendix

Table 13: Computational profile.

M

is the number of sampled lists and

T

the number of ReAct turns.

Extra lists $k$	Hit@10 (%)	Output tokens	Hit@10 (%)	Output tokens
	27B		35B
0	34.38	173	32.81	179
1	36.72	346	32.81	354
2	39.06	519	36.72	529
4	39.84	864	42.19	877
8	46.09	1,556	42.19	1,575

Appendix

Table 14: Sampling-budget ablation on 128 development cases with fixed model weights. Output tokens are mean totals per case across all

k+1

generations; they exclude input tokens, retrieval, and runtime overhead.

Explore similar work

RareDxR1: Autonomous Medical Reasoning for Rare Disease Diagnosis Beyond Human Annotation
Jun 30, 2026Deyang Jiang, Haoran Wu, Ziyi Wang +4Rare Disease Clinical Reasoning Training
The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, China · Peking Union Medical College Hospital, Beijing, China · School of Future Technology, University of Chinese Academy of Sciences, Beijing, China
A specialized reasoning large language model for accelerating rare disease diagnosis: a randomized AI physician assistance trial
Jun 23, 2026Haichao Chen, Songchi Zhou, Zhengyun Zhao +28Rare Disease Reasoning Skills
Tsinghua Medicine, Tsinghua University, Beijing, China. · Department of Ophthalmology, Peking Union Medical College Hospital, Chinese Academy of Medical Sciences & Peking Union Medical College, Beijing, China. · Department of Statistics and Data Science, Tsinghua University, Beijing, China. +19
Teaching agentic AI to generalize expert diagnostic reasoning in rare diseases
Jun 15, 2026Minh-Ha Nguyen, Erica Gray, Bryce A. Schuler +16Rare Disease Real-World Clinical Workflows
Department of Epidemiology, Vanderbilt University, Nashville, TN, USA · Department of Pediatrics, Vanderbilt University Medical Center, Nashville, TN, USA · Department of Medicine, Vanderbilt University Medical Center, Nashville, TN, USA +4

Jun 30, 2026, cs.AIJ/K move · Enter open · S save

RareDxR1: Autonomous Medical Reasoning for Rare Disease Diagnosis Beyond Human Annotation

Deyang Jiang, Haoran Wu, Ziyi Wang, Yiming Rong+3

The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, China · Peking Union Medical College Hospital, Beijing, China · School of Future Technology, University of Chinese Academy of Sciences, Beijing, China

Rare Disease Clinical Reasoning Training Human Annotations

Open paper AskPDF arXiv

Rare disease differential diagnosis is a critical yet arduous clinical task, requiring physicians to identify precise phenotypes from complex, unstructured patient symptoms and execute intricate reasoning within a vast search space. However, existing AI approaches typically rely on pipeline-based phenotype extraction or retrieval-augmented generation, which suffer from critical information loss due to predefined ontologies, retrieval bottlenecks, and a lack of diagnostic logic. To address these challenges, we introduce RareDxR1, an end-to-end reasoning-centric large language model designed for open-domain rare disease diagnosis directly from unstructured clinical notes. We design a progressive end-to-end training framework by synergizing knowledge internalization with autonomous evolutionary learning, thereby bypassing reliance on structured phenotypes and closed-set decision-making. To overcome the limitations of RAG and phenotype restriction, we enabled the deep internalization of fragmented rare-disease knowledge directly into the model's parameters. Moreover, to bridge the gap between model generation and expert reasoning, we propose Reflection-Enhanced Reasoning Sampling (RERS), a strategy that synthesizes expert-level diagnostic trajectories by learning from failures without human annotation. Additionally, we propose a dual-level curriculum reinforcement learning approach for gradually mastering rare disease diagnosis. Experimental results demonstrate that RareDxR1 achieves state-of-the-art accuracy across different benchmarks, marking a significant breakthrough in open-domain rare disease diagnosis. Our code and dataset will be publicly available.

Abstract

Explore similar work

RareDxR1: Autonomous Medical Reasoning for Rare Disease Diagnosis Beyond Human Annotation

A specialized reasoning large language model for accelerating rare disease diagnosis: a randomized AI physician assistance trial

Teaching agentic AI to generalize expert diagnostic reasoning in rare diseases