TRIO: Task Refinement with Individual Optimization for Retrieval-Augmented Generation

Yingtao Ren; Xiao Luo; Yiwei Fu; Jinzhao Zhou; Yu-Cheng Chang; Chin-Teng Lin

cs.CLNov 2, 2026

TRIO: Task Refinement with Individual Optimization for Retrieval-Augmented Generation

Authors: Yingtao Ren, Xiao Luo, Yiwei Fu, Jinzhao Zhou, Yu-Cheng Chang, Chin-Teng Lin

Organizations: University of Technology Sydney · University of Wisconsin–Madison · Peking University

Abstract

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]( $\downarrow$ 2.64% )	0.5555 [-0.3ex]( $\downarrow$ 3.73% )	0.4547 [-0.3ex]( $\downarrow$ 3.46% )	0.5191 [-0.3ex]( $\downarrow$ 8.29% )	0.1357 [-0.3ex]( $\downarrow$ 26.93% )	0.2456 [-0.3ex]( $\downarrow$ 12.85% )	0.4560 [-0.3ex]( $\downarrow$ 8.06% )	0.5590 [-0.3ex]( $\downarrow$ 9.84% )
TRIO w/o Judge Opt. [-0.3ex] (0)	0.4284 [-0.3ex]( $\downarrow$ 1.52% )	0.5587 [-0.3ex]( $\downarrow$ 3.17% )	0.4765 [-0.3ex]( $\uparrow$ 1.17% )	0.5416 [-0.3ex]( $\downarrow$ 4.31% )	0.1593 [-0.3ex]( $\downarrow$ 14.22% )	0.2635 [-0.3ex]( $\downarrow$ 6.49% )	0.4480 [-0.3ex]( $\downarrow$ 9.68% )	0.5643 [-0.3ex]( $\downarrow$ 8.98% )
TRIO w/o Answer Opt. [-0.3ex] (0)	0.3735 [-0.3ex]( $\downarrow$ 14.14% )	0.4926 [-0.3ex]( $\downarrow$ 14.63% )	0.2778 [-0.3ex]( $\downarrow$ 41.02% )	0.3618 [-0.3ex]( $\downarrow$ 36.08% )	0.1051 [-0.3ex]( $\downarrow$ 43.40% )	0.1855 [-0.3ex]( $\downarrow$ 34.17% )	0.2800 [-0.3ex]( $\downarrow$ 43.55% )	0.3967 [-0.3ex]( $\downarrow$ 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.

Panel of Figure 3 — 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:

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

Retrieval-augmented generation connects language models to external knowledge so answers can use current, domain-specific evidence.

See all 817 papers

KARE-RAG: Knowledge-Aware Refinement and Enhancement for RAG
Jun 3, 2025Yongjian Li, HaoCheng Chu, Yukun Yan +7Retrieval-Augmented Generation LLM Alignment
School of Biomedical Engineering, Tsinghua Medicine, Tsinghua University · 2Tsinghua Laboratory of Brain and Intelligence · School of Informatics, Xiamen University +6
Latent Abstraction for Retrieval-Augmented Generation
Apr 20, 2026Ha Lan N. T, Minh-Anh Nguyen, Dung D. LeRetrieval-Augmented Generation Adaptive RAG
Center for AI Research, VinUniversity, Vietnam
MASS-RAG: Multi-Agent Synthesis Retrieval-Augmented Generation
Apr 20, 2026Xingchen Xiao, Heyan Huang, Runheng Liu +1Retrieval-Augmented Generation Multi-Agent LLM Systems
School of Computer Science and Technology, Beijing Institute of Technology · Department of Mathematical Sciences, Tsinghua University

Jun 3, 2025J/K move · Enter open · S save

KARE-RAG: Knowledge-Aware Refinement and Enhancement for RAG

Yongjian Li, HaoCheng Chu, Yukun Yan, Zhenghao Liu+6

School of Biomedical Engineering, Tsinghua Medicine, Tsinghua University · 2Tsinghua Laboratory of Brain and Intelligence · School of Informatics, Xiamen University+6

Retrieval-Augmented Generation LLM Alignment

Open paper AskPDF arXiv

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.

Abstract

Retrieval-Augmented Generation

Explore similar work

KARE-RAG: Knowledge-Aware Refinement and Enhancement for RAG

Latent Abstraction for Retrieval-Augmented Generation

MASS-RAG: Multi-Agent Synthesis Retrieval-Augmented Generation