TAdaRAG: Task Adaptive Retrieval-Augmented Generation via On-the-Fly Knowledge Graph Construction

AI summary

Overview

Research area: Natural Language Processing — retrieval-augmented generation (RAG), knowledge-graph-augmented generation, long-context reasoning, and LLM fine-tuning (supervised fine-tuning plus reinforcement learning).

Technical level: Advanced. The paper assumes familiarity with RAG pipelines, chunking, knowledge graph extraction, LoRA fine-tuning, and policy-gradient (REINFORCE) optimization.

Scope: The paper proposes TAdaRAG, a framework that constructs task-adaptive knowledge graphs on the fly (rather than relying on pre-built graphs) and evaluates it on six public benchmarks plus a real-world business benchmark, NowNewsQA, using Mistral-7B-Instruct, Qwen2.5-7B-Instruct, and Qwen2.5-14B-Instruct.

What This Paper Is About

Traditional RAG splits retrieved documents into small chunks because of the input context window, which truncates complete knowledge, causes information loss, and leads to hallucinations and broken reasoning chains. It also feeds unorganized, unstructured text that contains irrelevant details. TAdaRAG's goal is to instead build a structured, domain-specific knowledge graph at inference time, so the model reasons over concise, logically organized, non-redundant knowledge tailored to the detected task intent.

Key Contributions

1. A task-adaptive RAG framework that integrates structured knowledge graph representations directly into the reasoning process. By dynamically constructing domain-relevant subgraphs, it targets hallucinations caused by chunked input in long-text tasks, strengthens reasoning in complex scenarios, and improves extraction of external knowledge.

2. An intent-driven routing plus two-stage training pipeline: intent detection routes input to a domain-specific extraction template; supervised fine-tuning teaches high-quality extraction; and a reinforcement-learning-based instruction-level implicit extraction mechanism (using REINFORCE) achieves automatic KG extraction and self-optimization.

3. Broad empirical evaluation on six public benchmarks (Health, Biology, Legal from ULTRADOMAIN; HotpotQA and 2WikiMQA from LongBench; GovReport) and a real-world business benchmark, NowNewsQA, across three backbone models, showing gains in factual QA, multi-hop reasoning, long-text summarization, and industry scenarios, plus human evaluation of answer quality.

4. Deployment evidence: TAdaRAG has been deployed in commercial applications, with trial accounts available for user access and testing. Code is released at https://github.com/IAAR-Shanghai/TAdaRAG.

Main Findings

• Hallucination mitigation (Mistral-7B-Instruct): TAdaRAG outperforms the SOTA baseline MEMORAG in factual domains — Health 37.40 → 40.77 and Biology 35.70 → 39.31. On Legal, TAdaRAG significantly outperforms NaïveRAG (35.80 → 49.88), though it remains close to but below MEMORAG (49.88 vs. 51.20).

• Reasoning enhancement: Compared to MEMORAG, TAdaRAG improves complex reasoning on 2WikiMQA (30.30 → 39.31) and multi-hop QA on HotpotQA (42.90 → 44.83).

• Task-oriented extraction: TAdaRAG outperforms MEMORAG on summarization (GovReport: 31.60 → 36.41, ROUGE-L).

• Consistent gains across backbones: On Qwen2.5-7B-Instruct, TAdaRAG reaches 42.38 (Health), 40.75 (Biology), 46.83 (Legal), 49.23 (HotpotQA), 43.79 (2WikiMQA), and 36.95 (GovReport), with asterisks marking statistically significant improvements (p < 0.01) over SOTA RAG baselines on Health, Biology, HotpotQA, 2WikiMQA, and GovReport — but not on Legal. On Qwen2.5-14B-Instruct, TAdaRAG reports 51.00, 51.32, 55.83, 56.28, 54.79, and 38.36 versus MEMORAG's 47.37, 42.33, 55.04, 52.91, 50.42, and 32.99.

• Ablation results: Introducing prompt-based KG (w/ graph) improves over NaïveRAG, with 2WikiMQA rising from 20.60 to 38.48 and GovReport from 27.40 to 33.72. Supervised fine-tuning (w/ sft) improves accuracy by 19.44% on Legal and 5.56% on Biology. Reinforcement learning (w/ reinforce) further improves Legal complex QA by 26.86% compared to Stage 1 and pushes results to SOTA across datasets.

• Long-context comparison: TAdaRAG is compared against Self-Extend, H2O+THINK, and SnapKV+THINK on the same base model across six datasets. It handles multi-document QA on HotpotQA and 2WikiMQA "rivaling dedicated long-context models without modifying storage," and excels on GovReport summarization.

• Business scenario evaluation: On NowNewsQA with Mistral-7B-Instruct, TAdaRAG achieves the highest overall average human-expert score (7.904 vs. 7.720), with notable leads in conciseness (8.251 vs. 7.637) and factuality (8.449 vs. 7.850).

• LLM-as-judge reliability: GPT-4o scores correlate strongly with human ratings, with Pearson coefficients of 0.706 (Relevance), 0.755 (Numerical Precision), 0.847 (Conciseness), 0.842 (Factuality), 0.828 (Timeliness), 0.850 (Comprehensiveness), 0.867 (Clarity), 0.925 (Coherence), and 0.881 (Insightfulness). Inter-rater agreement among three expert annotators on the 150 test samples is also consistently high across all models and nine dimensions.

• Hyperparameter sensitivity: Both Mistral-7B and Qwen2.5-7B perform best with 3 parallel subgraphs on most datasets. HotpotQA benefits from 2 subgraphs, while GovReport performs better with 4. Increasing beyond 3 (e.g., 4 or 5) may introduce noise for 7B-parameter models. Qwen experiments show stable performance across subgraph numbers, suggesting stronger models train more stably.

• Not reported here: The paper states that additional experiments — statistical significance testing (partially reported in Table 5), latency analysis, KG refinement evaluation, evidence verification, and case studies — appear in the extended version. Latency and cost numbers are therefore not available in this content.

Methodology in Plain English

The framework reformulates a generation task (e.g., summarization) as: given a document and its target answer, build a knowledge graph that best supports producing that answer.

Training stage 1 — supervised knowledge extraction fine-tuning. Because pretrained language models extract imprecise entities, the authors manually design extraction templates for general and specialized domains. Given a user query q and external knowledge r, the model uses prompts to detect intent and select the matching template t. Instructions of the form I = {q, r, t} are fed to strong LLMs to produce high-quality graphs G. This yields 9,548 fine-tuning samples spanning four question domains and seven sub-datasets. The model is then tuned with LoRA for 5 epochs (max input length 20,480 tokens, batch size 1, gradient accumulation over 8 steps, cosine learning rate starting at 5e−5).

Training stage 2 — task-adaptive KG construction. Since an LLM cannot judge the quality of a single graph in isolation, the model samples p parallel subgraphs per instruction, using learnable tokens <|startextraction|> and <|endextraction|> so the graph is embedded implicitly during generation. A "mixing network" concatenates the hidden state without the graph and with each subgraph, passing them through a three-layer MLP with ReLU to compute a per-token weight ω; the final log-likelihood is a weighted combination of the with-graph and without-graph log-likelihoods. A reward function based on REINFORCE gives positive reward only when a subgraph beats the average improvement over the no-graph baseline, and the total loss combines the base loss, graph loss, and REINFORCE loss with hyperparameters α and β. This stage uses ZeRO stage-2 with AdamW, per-GPU batch size 1, bfloat16 precision, 3 epochs at learning rate 5e−7, and sampling temperature T = 0.6; evaluation uses greedy decoding, and the maximum KG length is 2048 tokens. Total training takes roughly 16 hours on 8 NVIDIA A100 (80 GB) GPUs (4 hours for Stage 1, 12 hours for Stage 2).

Evaluation design. Six public datasets cover domains and task types (Health: 180 samples, Biology: 220, Legal: 438, HotpotQA: 200, 2WikiMQA: 200, GovReport: 200), scored with F1 except GovReport, which uses ROUGE-L. NowNewsQA contains 3,150 news-domain multi-document QA examples (3,000 training, 150 testing) retrieved with the Xinyu AI Search engine. Since gold-standard answers are difficult to define, evaluation uses nine multi-dimensional criteria — Relevance, Numerical Precision, Conciseness, Factuality, Timeliness, Comprehensiveness, Clarity, Coherence, and Insightfulness — rated by human journalism-domain experts and by GPT-4o at temperature 0. Seven baselines are compared: NaïveRAG, BGE-M3, RQ-RAG, GraphRAG, HippoRAG, MEMORAG, and PathRAG.

Why This Matters

Impact on research. The paper argues that graph-based RAG methods typically depend on pre-constructed or static KGs that need manual maintenance, lack scalability, and carry redundant or incomplete information. TAdaRAG's shift — building task-oriented graphs during reasoning rather than during retrieval — offers an alternative design point, and the release of code plus a production deployment gives the community a reproducible and industrially tested reference. The LLM-as-judge correlation study also adds evidence on when automated evaluation can substitute for costly human evaluation.

Real-world applications:

• News question answering: the NowNewsQA benchmark targets Chinese current-affairs news queries over redundant, noisy, partially relevant retrieved documents.
• Legal document analysis: the Legal results target answering questions that require integrating lengthy legal clauses and contract terminology.
• Biomedical and health QA: the Health and Biology datasets cover medical texts, diseases, drugs, treatments, cell biology, genetics, ecology, and evolutionary theory.
• Long-document summarization: GovReport involves U.S. Government Accountability Office and Congressional Research Service reports where key information is distributed across the entire document.

Industry relevance. The framework has been deployed in commercial applications with trial accounts available. It is co-authored with industry and engineering affiliations (MemTensor (Shanghai) Technology Co., Ltd.; China Haisum Engineering Co., Ltd.) and evaluated on documents mirroring actual search outputs from a production-level engine. The authors nonetheless note that dynamic KG construction and multi-stage training add computational overhead, and that reliance on manually crafted templates may constrain efficiency and adaptability.

Future Directions

1. Reduce computational cost and improve KG construction efficiency, since the current two-stage design adds overhead on top of normal inference and training cost.

2. Reduce reliance on manually crafted extraction templates, which currently require identifying key application domains and hand-selecting entity types — a constraint on adaptability in more complex or unseen scenarios.

3. Improve scalability and reliability in real-world settings, as stated in the conclusion's future work.

4. Extend the evaluation to the additional analyses promised in the extended version — latency analysis, KG refinement evaluation, evidence verification, and case studies — and investigate why the Legal domain shows TAdaRAG slightly below MEMORAG on two backbones despite otherwise consistent gains.

Target Audience

This paper is most useful to applied NLP researchers and engineers working on RAG systems, knowledge-graph-augmented generation, and long-context processing; to practitioners building enterprise QA or summarization products who need to decide between chunk-based RAG and graph-based alternatives; and to researchers interested in combining supervised fine-tuning with reinforcement learning to improve intermediate structured representations. Readers without background in RAG, graph extraction, or policy-gradient methods will find the methodology section demanding.

Authors’ abstract

Retrieval-Augmented Generation (RAG) improves large language models by retrieving external knowledge, often truncated into smaller chunks due to the input context window, which leads to information loss, resulting in response hallucinations and broken reasoning chains. Moreover, traditional RAG retrieves unstructured knowledge, introducing irrelevant details that hinder accurate reasoning. To address these issues, we propose TAdaRAG, a novel RAG framework for on-the-fly task-adaptive knowledge graph construction from external sources. Specifically, we design an intent-driven routing mechanism to a domain-specific extraction template, followed by supervised fine-tuning and a reinforcement learning-based implicit extraction mechanism, ensuring concise, coherent, and non-redundant knowledge integration. Evaluations on six public benchmarks and a real-world business benchmark (NowNewsQA) across three backbone models demonstrate that TAdaRAG outperforms existing methods across diverse domains and long-text tasks, highlighting its strong generalization and practical effectiveness.