TRACE: Deployable Tree-Relational Structure Enhancement for Oncology LLMs

Jizheng Lai; Yingyun Li; Ying Qin; Haiyang Qian

cs.CLSep 20, 2026

TRACE: Deployable Tree-Relational Structure Enhancement for Oncology LLMs

Authors: Jizheng Lai, Yingyun Li, Ying Qin, Haiyang Qian

Organizations: Pharmaron, China · AI Starfish, China

Abstract

Large language models are increasingly used in oncology applications, but their predictions are often weakly grounded in explicit medical structure. We present TRACE, a deployable tree-relational enhancement framework for oncology LLMs. TRACE separates expensive offline structure learning from lightweight online inference: oncology concepts and relations are organized into an updatable tree-relational structure, refined using LM-loss-derived evidence, and retrieved at inference time as compact prompt evidence. This design supports task-adaptive evidence selection without requiring supervised labels in the zero-shot setting. Across ten oncology classification tasks and one MedQuAD CancerGov QA benchmark, TRACE improves both label-free evaluation and supervised fine-tuning. Additional analyses show that TRACE improves over vanilla RAG and generic GraphRAG, remains useful under leakage-controlled METABRIC inputs, and produces interpretable evidence paths aligned with clinical reasoning. These results suggest that explicit, updatable medical structure is a practical path toward more accurate and auditable oncology LLM deployment.

Figures & tables

Figure 1: Overview of TRACE. Offline, TRACE retrieves concept-query corpus samples, performs structure-oriented LoRA adaptation, estimates LM-loss evidence, and updates tree nodes and graph edges with Bayesian rules. Online, TRACE retrieves compact tree-relational evidence and injects it into downstream prompts. Supervised settings can additionally learn task-branch weights.

Task	Baichuan- M2-32B	Baichuan- M3- 235B	Llama-4- Scout	HuatuoGPT- 3-32B	Qwen3- 32B	DeepSeek- V4- Flash	Qwen3.5- 27B	TRACE
Oncology examination QA
Professional Practice	0.785	0.840	0.774	0.796	0.807	0.840	0.818	0.906
Oncology Attending	0.727	0.787	0.628	0.740	0.749	0.815	0.781	0.807
Related Knowledge	0.792	0.833	0.657	0.790	0.794	0.853	0.836	0.845
Private clinical records
Common Cancer Dx.	0.701	0.733	0.649	0.710	0.718	0.725	0.718	0.725

Table 1: Label-free zero-shot accuracy on ten oncology classification tasks. TRACE uses Qwen3.5-27B as the backbone and injects retrieved tree-relational evidence without supervised training on target-task labels. All baselines are evaluated under the same label-free setting. Each dataset block reports per-task accuracy and a group-level macro average, and the last row reports the macro average over all ten tasks. Bold and underline denote the best and second-best results in each row.

Model / Method	Overall Acc.	Overall F1	Task-macro Acc.	Task-macro F1
Qwen3-32B LoRA	0.755	0.389	0.758	0.546
HuatuoGPT-3-32B LoRA	0.748	0.364	0.751	0.526
Baichuan-M2-32B LoRA	0.744	0.359	0.744	0.526
Qwen3.5-27B LoRA	0.757	0.361	0.760	0.528
Llama-4-Scout LoRA	0.758	0.360	0.761	0.528
TRACE	0.797	0.480	0.797	0.579

Table 2: Fine-tuning results on the five METABRIC tasks. Overall metrics are computed by pooling all test examples, while task-macro metrics average the corresponding metric across the five tasks. Bold denotes the best result in each column.

Method	Exam QA	Clinical records	METABRIC	10-task Avg.
Qwen3.5-27B	0.811	0.604	0.728	0.728
+ Concept Query LoRA	0.816	0.610	0.714	0.723
+ Tree-Relational RAG	0.847	0.615	0.736	0.745
TRACE	0.853	0.622	0.735	0.748
Vanilla RAG	0.817	0.623	0.727	0.733
Generic GraphRAG	0.828	0.624	0.725	0.736

Table 3: Main ablation summary. All numbers are group-level or 10-task macro accuracy. Bold denotes the best result within each comparison block.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Stage	Input	Output	GPUs	Wall-clock	GPU-hours
Initial structure construction	Answer-blind task/corpus text and oncology literature	300 concepts / 280 relations	N/A	Not separately logged	N/A
Retrieval and update-data preparation	300 concepts / 280 relations	LM-loss training/scoring corpus	2	3.3 min	0.11
LoRA warmup training	LM-loss corpus	Structure-adapted LoRA adapter	2	87.8 min	2.93
LM-loss scoring and Bayesian update	Initial tree/graph and LoRA adapter	Updated tree/graph weights	1	6.1 min	0.10
Total refinement	300 concepts / 280 relations	Refined structure	–	97 min	3.1–3.3

Appendix

Table 4: Deployment cost of structure construction and one LM-loss Bayesian refinement round. The initial structure is reported by scale and provenance because its original construction logs were not retained as separately profiled GPU jobs. GPU-hours are computed from measured wall-clock time and the number of GPUs used in each profiled stage.

Metric	Median	P90	P95	Mean
Retrieval latency (ms)	2.71	5.54	7.43	3.43
Base latency (ms)	304.99	439.31	517.84	338.01
TRACE latency (ms)	404.24	619.67	669.72	447.14
TRACE overhead (ms)	101.50	137.89	236.90	109.12
Base prompt tokens	147.5	630.4	848.3	328.9
TRACE prompt tokens	795.5	1230.2	1438.6	942.7

Appendix

Table 5: Inference latency of TRACE under direct HuggingFace Transformers + PEFT inference. Statistics are computed over 100 samples from the three examination QA tasks and two clinical-record tasks, with 20 samples per task, using Qwen3.5-27B with the LoRA adapter on a single GPU and

\texttt{max\_new\_tokens}=1

. TRACE latency includes structural retrieval and prompt-augmented model inference.

Dataset	# examples	Description
Oncology examination QA
Professional Practice	181	Chinese examination-style oncology questions.
Oncology Attending Physician	1,581	Chinese oncology attending-physician examination questions.
Related Professional Knowledge	708	Chinese medical professional-knowledge examination questions.
Private de-identified clinical records
Common Cancer Diagnosis	131	De-identified clinical records labeled with major cancer type.

Appendix

Table 6: Dataset summary. METABRIC tasks are derived from cBioPortal brca_metabric ; CancerGov QA is derived from the MedQuAD CancerGov subset.

Setting	Input fields
Default METABRIC	All non-target structured fields, including clinical, pathological, molecular, treatment, and outcome fields.
Leakage-controlled	Target field removed as before; treatment and outcome fields are additionally removed, including endocrine therapy, chemotherapy, radiotherapy, surgery, recurrence, survival, follow-up, and death-related fields.

Appendix

Table 7: Input-field control for METABRIC-derived tasks.

Model / Method	ROUGE-L	BERTScore Tiny	BERTScore RoBERTa
Baichuan-M2-32B LoRA	0.175	0.809	0.847
HuatuoGPT-3-32B LoRA	0.135	0.791	0.835
Qwen3-32B LoRA	0.155	0.805	0.845
Qwen3.5-27B LoRA	0.254	0.828	0.872
Llama-4-Scout LoRA	0.201	0.823	0.860
TRACE	0.278	0.848	0.884

Appendix

Table 8: Fine-tuning results on MedQuAD CancerGov QA generation. Models are evaluated by ROUGE-L and BERTScore. Bold denotes the best result in each column.

Noise	Method	Acc.	Macro-F1	Retention	Tokens	$\Delta$ Latency (ms)
Clean	LoRA-only	0.6779	0.5033	1.0000	392.6	0.0
Clean	Tree-Relational RAG	0.6840	0.5030	1.0000	644.3	+590.6
Clean	TRACE	0.6918	0.5063	1.0000	990.2	+1420.9
Moderate	LoRA-only	0.6686	0.4914	0.9899	390.3	0.0
Moderate	Tree-Relational RAG	0.6775	0.4939	0.9926	643.5	+595.6
Moderate	TRACE	0.6921	0.5101	1.0022	990.2	+1425.7

Appendix

Table 9: Robustness under EMR-like input noise. Retention is computed relative to each method’s clean accuracy. Latency overhead is measured relative to LoRA-only inference under the same noise condition.

Top- $k$	Acc.	Macro-F1	Tokens	Latency (ms)
0	0.7213	0.5296	186.6	698.2
2	0.7282	0.5278	355.7	1184.1
4	0.7313	0.5317	494.7	1592.3
8	0.7371	0.5330	709.7	2231.2
12	0.7414	0.5393	819.2	2547.3
16	0.7449	0.5408	871.1	2701.7

Appendix

Table 10: Accuracy-latency tradeoff as the number of retrieved structural evidence items increases. Larger evidence budgets improve accuracy but increase prompt length and inference latency.

Model / Method	Exam QA	Clinical records	METABRIC	10-task Macro Avg.
Baichuan-M2-32B	0.725	0.370	0.463	0.523
Baichuan-M3-235B	0.774	0.396	0.453	0.538
Llama-4-Scout	0.677	0.331	0.326	0.432
HuatuoGPT-3-32B	0.731	0.463	0.441	0.532
Qwen3-32B	0.739	0.469	0.452	0.542
DeepSeek-V4-Flash	0.790	0.474	0.480	0.572

Appendix

Table 11: Group-level task-macro F1 for label-free zero-shot evaluation. Each group column averages macro-F1 over tasks in the corresponding dataset group, and the last column averages macro-F1 over all ten tasks. Bold denotes the best result in each column.

Method	Exam QA Acc. (95% CI)	Clinical records Acc. (95% CI)	METABRIC Acc. (95% CI)	10-task Macro Acc. (95% CI)
Baichuan-M2-32B	0.768 [0.744, 0.791]	0.582 [0.533, 0.630]	0.692 [0.683, 0.700]	0.693 [0.680, 0.705]
Baichuan-M3-235B	0.820 [0.798, 0.840]	0.617 [0.568, 0.664]	0.697 [0.688, 0.706]	0.718 [0.706, 0.730]
Llama-4-Scout	0.686 [0.661, 0.710]	0.461 [0.413, 0.509]	0.640 [0.631, 0.649]	0.618 [0.605, 0.631]
HuatuoGPT-3-32B	0.775 [0.751, 0.798]	0.604 [0.555, 0.652]	0.668 [0.659, 0.676]	0.687 [0.674, 0.700]
Qwen3-32B	0.783 [0.760, 0.805]	0.610 [0.561, 0.658]	0.673 [0.665, 0.682]	0.694 [0.681, 0.706]
DeepSeek-V4-Flash	0.836 [0.815, 0.856]	0.615 [0.567, 0.663]	0.688 [0.679, 0.697]	0.718 [0.706, 0.730]

Appendix

Table 12: Bootstrap confidence intervals for the main zero-shot accuracy results. Intervals are computed with 10,000 task-stratified bootstrap samples. Each group score is the macro average over tasks in that group, and the final column is the macro average over all ten tasks. Bold denotes the best point estimate in each column.

Task	Qwen3.5 -27B	TRACE	$\Delta$
Professional Practice	0.818	0.906	+0.088
Oncology Attending	0.781	0.807	+0.026
Related Knowledge	0.836	0.845	+0.009
Common Cancer Dx.	0.718	0.725	+0.007
Breast Stage Pred.	0.491	0.519	+0.028
METABRIC Stage	0.729	0.718	-0.011

Appendix

Table 13: Per-task zero-shot accuracy gain of TRACE over Qwen3.5-27B.

Evidence	Prof. Practice	Oncology Attending	Related Knowledge	Common Cancer Dx.	Breast Stage Pred.	5-task Macro Avg.
Base	0.818	0.781	0.836	0.718	0.491	0.729
Tree only	0.823	0.791	0.838	0.725	0.516	0.739
Graph only	0.879	0.807	0.843	0.710	0.547	0.757
Tree+Graph	0.906	0.807	0.845	0.725	0.519	0.760

Appendix

Table 14: Tree and graph evidence ablation. All numbers are accuracy. The table compares tree-only, graph-only, and tree-plus-graph evidence, and reports per-task accuracy with a macro average over the five evaluated tasks. Bold denotes the best result in each column.

Update rule	Exam QA	Clinical records	5-task Macro Avg.
Initial weight	0.838	0.624	0.753
Linear update	0.841	0.615	0.751
Beta-Bernoulli	0.843	0.617	0.752

Appendix

Table 15: Accuracy under different structure-weight update rules. The results are similar across update mechanisms; Beta-Bernoulli is used for its bounded probabilistic interpretation and support for iterative refinement, rather than for a consistent accuracy advantage.

Model	Variant	Oncology examination QA	Private clinical records	METABRIC	10-task Macro Avg.
Baichuan-M2-32B	Base	0.768	0.582	0.692	0.693
Baichuan-M2-32B	+ Tree-Relational RAG	0.779 (+0.012)	0.578 (-0.004)	0.705 (+0.013)	0.702 (+0.009)
Baichuan-M3-235B	Base	0.820	0.617	0.697	0.718
Baichuan-M3-235B	+ Tree-Relational RAG	0.845 (+0.025)	0.601 (-0.016)	0.705 (+0.008)	0.726 (+0.008)
Llama-4-Scout	Base	0.686	0.461	0.640	0.618
Llama-4-Scout	+ Tree-Relational RAG	0.785 (+0.099)	0.559 (+0.098)	0.647 (+0.007)	0.671 (+0.053)

Appendix

Table 16: Backbone-level ablation of Tree-Relational RAG. Each model is compared with and without the same learned structural retrieval module. Numbers in parentheses indicate the change over the corresponding base model.

Method	Overall Acc.	Overall F1	Task-macro Acc.	Task-macro F1
TRACE w/o task-branch weighting	0.7614	0.3853	0.7628	0.5376
TRACE	0.7969	0.4802	0.7968	0.5792

Appendix

Table 17: Task-branch weighting ablation on METABRIC mixed fine-tuning. Metrics follow the same definition as Table 2 .

Method	Prof. Practice	Oncology Attending	Related Knowledge	Common Cancer Dx.	Breast Stage Pred.	Task-macro
Vanilla RAG	0.1380	0.7038	0.8624	0.4729	0.4546	0.5263
Generic GraphRAG	0.1365	0.7286	0.8815	0.4915	0.4495	0.5375

Appendix

Table 18: Macro-F1 comparison between retrieval baselines on the three examination QA tasks and two clinical-record tasks. Both baselines use the same oncology corpus as external evidence. Bold denotes the best result in each column.

Method	Stage	Histology	ER	PR	Relapse	Avg.
Accuracy
Qwen3.5-27B	0.662	0.762	0.755	0.744	0.594	0.704
Vanilla RAG	0.656	0.762	0.758	0.745	0.611	0.707
Generic GraphRAG	0.671	0.762	0.737	0.744	0.611	0.705
TRACE	0.741	0.763	0.777	0.744	0.620	0.729
Macro-F1

Appendix

Table 19: Leakage-controlled METABRIC evaluation after removing treatment-related and outcome-related input fields. All methods are evaluated on the same filtered inputs. Avg. reports task-macro average over the five METABRIC-derived tasks. Bold denotes the best result in each column; ties after rounding are left unmarked.

Panel of Figure 2 — Figure 2: Structural evidence case studies. TRACE retrieves compact task-relevant substructures, including local tree evidence and key graph relations, that connect input clues to the final prediction.

Medical LLMs

Medical large language models are AI systems adapted to clinical and biomedical data for health reasoning, information extraction, and communication.

See all 397 papers

OA-RAG: Ontology-based Agentic RAG for Faithful LLM Reasoning
Oct 26, 2026Hui Yang, Jiaoyan Chen, Yongsheng GaoAgentic RAG Retrieval-Augmented Generation
The University of Manchester, Manchester, United Kingdom · SNOMED International, London, United Kingdom
MedJudgeRAG: Option-Wise Evidence Judgment with Dynamic Knowledge Graphs for Medical MCQA
Jul 24, 2026Seongwon Seo, Seung Hwan Cho, Young-Min KimGraph Retrieval-Augmented Generation Medical QA
Department of Industrial Data Engineering, Hanyang University, Seoul, South Korea · School of Interdisciplinary Industrial Studies, Hanyang University, Seoul, South Korea
When Cases Get Rare: A Retrieval Benchmark for Off-Guideline Clinical Question Answering
May 20, 2026Doeun Lee, Muge Zhang, Yi Yu +11LLM Grounding Healthcare
The Ohio State University · The Ohio State University Wexner Medical Center · University of Chicago Medical Center

Oct 26, 2026J/K move · Enter open · S save

OA-RAG: Ontology-based Agentic RAG for Faithful LLM Reasoning

Hui Yang, Jiaoyan Chen, Yongsheng Gao

The University of Manchester, Manchester, United Kingdom · SNOMED International, London, United Kingdom

Agentic RAG Retrieval-Augmented Generation LLM Grounding Ontology-Based Reasoning Medical LLMs

Open paper AskPDF arXiv

Large language models (LLMs) perform well across many real-world tasks but remain prone to hallucination, limiting their reliability in domains such as biomedicine and healthcare, where answers must be grounded in established knowledge. Retrieval-augmented generation (RAG) mitigates this problem by incorporating external information, but often overlooks logical relationships and relies heavily on the LLM's own reasoning, which can compromise answer faithfulness. We propose OA-RAG, an agent that integrates formal reasoning with LLM reasoning, drawing on both the logical structure and informal annotations of domain ontologies. For each answer, OA-RAG reports the ontology observations used as evidence and a model-assessed grounding level (entailed, supported, related, or unrelated) that characterises the strength of that evidence. Across benchmarks derived from three large-scale ontologies and six real-world medical question-answering tasks, OA-RAG achieves the best overall performance among the evaluated ontology-fine-tuned and popular RAG methods, with particularly strong results on reasoning-intensive tasks. The reported grounding levels also offer a useful indication of answer reliability, with stronger grounding generally associated with higher accuracy. Finally, our results demonstrate the value of ontologies as knowledge sources: OA-RAG is competitive with, and sometimes outperforms, domain-specific RAG systems that draw on sources such as textbooks.

Abstract

Medical LLMs

Explore similar work

OA-RAG: Ontology-based Agentic RAG for Faithful LLM Reasoning

MedJudgeRAG: Option-Wise Evidence Judgment with Dynamic Knowledge Graphs for Medical MCQA

When Cases Get Rare: A Retrieval Benchmark for Off-Guideline Clinical Question Answering