This paper studies retrieval-augmented generation (RAG), which leverages external knowledge to improve the performance of language models. Despite the remarkable progress, current RAG approaches are still far from satisfactory due to inadequate retrieval and potential hallucinations. Towards this end, we propose a novel approach named Task Refinement with Individual Optimization (TRIO) for RAG. The core idea of our TRIO is to equip three well-designed agents with an iterative refinement paradigm by aligning their individual optimization objectives. Specifically, we first expand ambiguous queries with a query extension agent and quantify the reward based on the improved answer quality, which can be utilized to refine our agent for higher expansion efficiency. To enhance the retrieval quality, we include a unanimity voting to annotate the expanded query as insufficient or sufficient, and their generative outcomes are utilized as ground truth to optimize the judge agent. To further improve evidence utilization and answer accuracy, an answer agent is optimized with dynamic matching-based rewards with curriculum learning for final outputs. Extensive experiments across multiple benchmark datasets validate the effectiveness of the proposed TRIO in comparison with state-of-the-art approaches. Our code is available at https://github.com/yingtaoren/TRIO.
Figures & tables
Figure 1: Problem illustration. RAG struggles with inefficient knowledge acquisition in the retrieval stage and knowledge underutilization in the answer stage.
Figure 2: An overview of our TRIO framework. The tight cooperation among three specific optimized agents improves the retrieval efficiency and answer quality.
In-Domain
Out-of-Domain
Method
HotpotQA
2Wiki
NQ
Musique
Bamboogle
AVG
Single-Hop Retrieval
Qwen3-4B
0.2770
0.1584
0.3357
0.0571
0.2240
0.2074
Llama3.1-8B
0.3095
0.1639
0.3720
0.0533
0.1600
0.2237
Qwen3-30B
0.2979
0.2479
0.3175
0.1278
0.2400
0.2606
Multi-Hop Retrieval
Table 1: The overall evaluation results of TRIO and other baselines on five benchmarks. Red marks the best-performing method, and blue represents the second-best-performing method. All methods are evaluated under the same settings. Our TRIO achieves the strongest performance.
Hotpotqa
2Wiki
Musique
Bamboogle
Variants
EM
F1
EM
F1
EM
F1
EM
F1
TRIO w/o Query Opt. [-0.3ex] (0)
0.4235 [-0.3ex]( ↓ 2.64% )
0.5555 [-0.3ex]( ↓ 3.73% )
0.4547 [-0.3ex]( ↓ 3.46% )
0.5191 [-0.3ex]( ↓ 8.29% )
0.1357 [-0.3ex]( ↓ 26.93% )
0.2456 [-0.3ex]( ↓ 12.85% )
0.4560 [-0.3ex]( ↓ 8.06% )
0.5590 [-0.3ex]( ↓ 9.84% )
TRIO w/o Judge Opt. [-0.3ex] (0)
0.4284 [-0.3ex]( ↓ 1.52% )
0.5587 [-0.3ex]( ↓ 3.17% )
0.4765 [-0.3ex]( ↑ 1.17% )
0.5416 [-0.3ex]( ↓ 4.31% )
0.1593 [-0.3ex]( ↓ 14.22% )
0.2635 [-0.3ex]( ↓ 6.49% )
0.4480 [-0.3ex]( ↓ 9.68% )
0.5643 [-0.3ex]( ↓ 8.98% )
TRIO w/o Answer Opt. [-0.3ex] (0)
0.3735 [-0.3ex]( ↓ 14.14% )
0.4926 [-0.3ex]( ↓ 14.63% )
0.2778 [-0.3ex]( ↓ 41.02% )
0.3618 [-0.3ex]( ↓ 36.08% )
0.1051 [-0.3ex]( ↓ 43.40% )
0.1855 [-0.3ex]( ↓ 34.17% )
0.2800 [-0.3ex]( ↓ 43.55% )
0.3967 [-0.3ex]( ↓ 36.02% )
TRIO (Full)
0.4350
0.5770
0.4710
0.5660
0.1857
0.2818
0.4960
0.6200
Table 2: Ablation study results on the four multi-hop benchmarks with three variants.
Figure 3: (a) and (b): Visualization of EM and F1 metrics for TRIO with single or mixed reward training. (c): The reward values trajectory is shown. Mixed reward design surpasses single reward design in overall performance.
Musique
Bamboogle
Method
EM
F1
EM
F1
Fixed Round = 1
0.1556
0.2597
0.4480
0.5614
Fixed Round = 2
0.1585
0.2609
0.4480
0.5587
Fixed Round = 3
0.1585
0.2599
0.4560
0.5811
Fixed Round = 4
0.1547
0.2561
0.4560
0.5816
Fixed Round = 5
0.1531
0.2516
0.4560
0.5818
Table 3: TRIO performance with flexible and different fixed iteration rounds on the retrieval step.
Inference Time (s / 100 queries)
Method
Hotpotqa
2Wiki
Musique
Bamboogle
Avg
SuRe
603.2
711.8
521.7
687.4
631.0
IRCOT
64.8
90.1
84.5
82.7
80.5
Iter-RetGen
20.4
23.2
20.5
20.4
21.1
Self-RAG
10.9
10.2
9.6
10.1
10.2
Search-R1
91.5
86.4
97.4
85.1
90.1
Table 4: Inference time of eight methods (seconds per 100 queries) across different datasets.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Prompt for Query Extension Agent
System Prompt:
You are a retrieval query generation agent.
Input:
QUESTION: …
EXISTING QUERIES and RETRIEVED DOCUMENTS: …
Appendix
Table 8
Prompt for Answer Agent
System Prompt:
Answer the question based on the given document. Output only the final answer with no explanations or additional text.
Query Pool:
|Current Query Pool|
User Prompt:
Question: |Input Question|
Appendix
Table 9
Prompt for Judge Agent
System Prompt:
Determine whether the EXISTING QUERIES and RETRIEVED DOCUMENTS provide sufficient information to answer the QUESTION correctly. Output only one of the following: ‘Sufficient Information’ or ‘Insufficient Information’. Do not include explanations or additional text.
User Prompt:
Question: | Input Question |
Query Pool: | Current Query Pool |
Appendix
Table 10
Prompt for Judge Agent in Training
System Prompt:
Determine whether the EXISTING QUERIES and RETRIEVED DOCUMENTS provide sufficient information to answer the QUESTION correctly. Output only one of the following: ‘Sufficient Information’ or ‘Insufficient Information’. Do not include explanations or additional text.
User Prompt:
Question: Are the directors of both films I Can Do Bad All By Myself (Film) and Shit Year from the same country?
EXISTING QUERIES and RETRIEVED DOCUMENTS
Query 1: Are the directors of both films I Can Do Bad All By Myself (Film) and Shit Year from the same country?
Appendix
Table 11
Prompt for Query Extension Agent in Training
System Prompt:
You are a retrieval query generation agent.
Input:
QUESTION: …
EXISTING QUERIES and RETRIEVED DOCUMENTS: …
Appendix
Table 12
Prompt for Answer Agent in Training
System Prompt:
Answer the question based on the given document. Output only the final answer with no explanations or additional text.
EXISTING QUERIES and RETRIEVED DOCUMENTS
Query 1: Are the directors of both films I Can Do Bad All By Myself (Film) and Shit Year from the same country?
Retrieved Content:
Doc 1(Title: "I Can Do Bad All by Myself (film) "): …
Appendix
Table 13
Prompt for Moderator in Retrieval Debate
System Prompt:
You are the judge in a debate. Your task is to evaluate the arguments from agents. There are two types of agents:
1. Proponent Agent: Argues that the current retrieved content is sufficient.
2. Opponent Agent: Argues that the current retrieved content is insufficient and proposes query refinement.
User Prompt:
Appendix
Table 14
Prompt for Proponent in Retrieval Debate
System Prompt:
You are a debater. Argue that the current retrieved content is sufficient to answer the question and try to give the answer based on the given documents. Deliver a brief, strong argument with clear reasoning. Do not suggest further retrieval. No extra explanations.
User Prompt:
Question: |Input Question|
Query Pool: |Current Query Pool|
Appendix
Table 15
Prompt for Opponent in Retrieval Debate
System Prompt:
You are a critical thinker and debater. Your task is to challenge the sufficiency of the current documents. However, your ultimate goal is to find the correct answer efficiently. If you believe the provided information is TRULY and COMPLETELY sufficient, your duty is to concede.
The action you can choose:
1. Query Expansion: If a completely new line of inquiry is needed. Output exactly in this format in the end of your response: Query Expansion: [New Query]
2. Concede: If the current information is truly sufficient. Only Output: Concede
Appendix
Table 16
Case 1: Immediate Success (Early Stop)
Question: Who played the character in the Santa Clause 3 that has a series named after it that includes Frost at Christmas?
Ground Truth: Martin Short
Round 0:
Retrieval Query: Who played the character in the Santa Clause 3 that has a series named after it that includes Frost at Christmas?
Retrieved Content: 1. The Santa Clause 3 ; 2. The Santa Clause (film series) ; 3. The Santa Clause 3: The Escape Clause …
Judge: Sufficient (Decision: Stop & Answer)
Appendix
Table 5: Case studies demonstrating the token efficiency of TRIO. Driven by the coordination among the three agents, the system completes tasks with highly concise token outputs.
Agent
Stage
GPU-hours (H100)
TRIO-8B (Llama-8B)
Answer agent
Stage 1
33.38h
Answer agent
Stage 2
13.06h
Query extension agent
–
33.96h
Judge agent
–
12.60h
Total
93.00h
Appendix
Table 6: Training resource consumption of TRIO-8B and TRIO-4B.
MuSiQue
Bamboogle
Method
Low
High
Low
High
SuRe
0.0645
0.0508
0.1280
0.1200
IRCoT
0.0662
0.0559
0.1520
0.1920
Iter-RetGen
0.0699
0.0488
0.1760
0.1600
DRAG
0.1129
0.1067
0.2720
0.2640
Self-RAG
0.0384
0.0372
0.0640
0.0560
Appendix
Table 7: EM scores under low and high noise settings. Low and High denote replacing one and three of the top-5 retrieved chunks with noisy chunks, respectively.
Method
HotpotQA
2Wiki
MuSiQue
Bamboogle
SuRe
0.3538
0.1917
0.1296
0.2727
IRCoT
0.2825
0.2363
0.1244
0.3112
Iter-RetGen
0.4419
0.2636
0.1297
0.2701
DRAG
0.4146
0.2791
0.2010
0.4030
Self-RAG
0.2842
0.2351
0.1195
0.1576
Search-R1
0.5306
0.4879
0.2608
0.5780
Appendix
Table 8: F1 results on four multi-hop benchmarks. The best result in each column is shown in bold, and the second-best result is underlined.
Method
PubMedQA
SuRe
0.256
IRCoT
0.458
Iter-RetGen
0.716
DRAG
0.644
Self-RAG
0.628
Search-R1
0.568
Appendix
Table 9: Results of TRIO and other baselines on PubMedQA benchmark.
Retrieval-Augmented Generation (RAG) equips large language models with external knowledge and is central to knowledge-intensive tasks. As RAG systems enter real-world use, generators must reliably leverage retrieved evidence. Recent fine-tuning methods improve adaptation to RAG scenarios, but optimization remains challenging because retrieval may return incomplete, fragmented, noisy, or conflicting contexts. Complex tasks further require fine-grained evidence dependencies. These challenges make high-quality supervision costly and limit generalization. We present KARE-RAG (Knowledge-Aware Refinement and Enhancement for RAG), a training-time scaffolded alignment framework. It uses structured knowledge representations as temporary scaffolds to expose evidence organization, support localized factual refinement, and construct fine-grained preference pairs. In our main implementation, an expert LLM refines a lightweight graph-structured evidence sketch. The generator is optimized with token-weighted Dense Direct Preference Optimization (DDPO), which focuses learning on edited scaffold regions. Scaffolds are used only for data construction and training supervision. At inference time, the model runs as standard Vanilla RAG without graph construction, extra retrieval, or latency overhead. Experiments show that KARE improves transfer across the evaluated QA and relation extraction datasets with limited training data, while leaving general capabilities largely unchanged. KARE can also complement existing RAG training objectives as an additional alignment stage.
Yongjian Li, HaoCheng Chu, Yukun Yan +7
School of Biomedical Engineering, Tsinghua Medicine, Tsinghua University · 2Tsinghua Laboratory of Brain and Intelligence · School of Informatics, Xiamen University +6
Retrieval-Augmented Generation (RAG) has become a standard approach for enhancing large language models (LLMs) with external knowledge, mitigating hallucinations, and improving factuality. However, existing systems rely on generating natural language queries at each hop and maintaining a strict architectural separation between retriever and generator, preventing them from leveraging the full representational capacity of the LLM. We propose LAnR (Latent Abstraction for RAG), a unified framework in which a single LLM jointly performs encoding, retrieval, and generation entirely within its own latent space. Rather than generating textual queries, LAnR produces dense retrieval vectors from the hidden states of a designated [PRED] token and uses them to match against encoded document representations from the same model. Furthermore, LAnR adaptively decides when sufficient evidence has been retrieved using a lightweight MLP control head over those same hidden states, eliminating both the separate retriever and explicit token-level stopping reasoning. This design is motivated by our empirical observation that answer token entropy reliably signals retrieval sufficiency. Extensive experiments on six QA benchmarks spanning single-hop and multi-hop settings demonstrate that LAnR outperforms existing RAG methods, while achieving improved inference efficiency through reduced number of retrieval calls and tighter model integration.
Large language models (LLMs) are widely used in retrieval-augmented generation (RAG) to incorporate external knowledge at inference time. However, when retrieved contexts are noisy, incomplete, or heterogeneous, a single generation process often struggles to reconcile evidence effectively. We propose MASS-RAG, a multi-agent synthesis approach to retrieval-augmented generation that structures evidence processing into multiple role-specialized agents. MASS-RAG applies distinct agents for evidence summarization, evidence extraction, and reasoning over retrieved documents, and combines their outputs through a dedicated synthesis stage to produce the final answer. This design exposes multiple intermediate evidence views, allowing the model to compare and integrate complementary information before answer generation. Experiments on four benchmarks show that MASS-RAG consistently improves performance over strong RAG baselines, particularly in settings where relevant evidence is distributed across retrieved contexts.
Xingchen Xiao, Heyan Huang, Runheng Liu +1
School of Computer Science and Technology, Beijing Institute of Technology · Department of Mathematical Sciences, Tsinghua University