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

AI summary

Overview

Research area: Natural Language Processing — retrieval-augmented generation (RAG) and multi-agent LLM systems. Technical level: Intermediate to Advanced (assumes familiarity with RAG pipelines and reinforcement learning for LLMs). Scope: The paper proposes TRIO, a three-agent RAG framework in which a query extension agent, a judge agent, and an answer agent are each separately trained and then coordinated in an iterative retrieval loop, evaluated on five QA benchmarks.

What This Paper Is About

RAG systems try to improve language-model answers by retrieving external documents, but they often retrieve the wrong things and then fail to use what they retrieved, which produces hallucinations. TRIO attacks both halves of that problem by splitting the pipeline into three roles — writing better retrieval queries, deciding when enough evidence has been gathered, and answering from the evidence — and giving each role its own training signal. The goal is a system where agents are specialized rather than sharing one backbone, and where their objectives are aligned so the loop works as a whole.

Key Contributions

1. Role-specialized agent decomposition of RAG. TRIO assigns three distinct subtasks to three separately optimized agents — query extension, sufficiency judgment, and answer generation — instead of using one uniform backbone across all stages, as prior multi-agent RAG work (e.g., DRAG) does.
2. Data-generation and training protocol that ties the agents together. Starting from question–answer supervision only, the authors design distinct data construction methods for each subtask, so that the query agent is trained to retrieve evidence aligned with the answer agent's preferences and the judge is trained to understand the other agents' behavior.
3. Per-agent optimization strategies. Two-stage outcome-based reinforcement learning (GRPO) with a mixed reward for the answer agent, an outcome-based reward defined as the change in answer F1 for the query agent, and inductive supervised fine-tuning with unanimity-voting labels for the judge agent.
4. Empirical validation across five benchmarks plus ablations, retrieval-round analysis, reward-design analysis, and inference-time measurements, showing gains over baselines including Search-R1, R1-Searcher, and DRAG.

Main Findings

• Best average EM among compared methods. TRIO-8B (Llama3.1-8B-Instruct backbone) reaches an average exact-match score of 0.4294, versus 0.4034 for Search-R1, 0.3779 for R1-Searcher, and 0.2490 for DRAG. The paper states TRIO-8B improves 6.45% on average over the second-best method and 72.45% over DRAG.
• Four of five benchmarks won. TRIO-8B scores 0.4350 on HotpotQA, 0.4710 on 2Wiki, 0.1857 on Musique, and 0.4960 on Bamboogle. Search-R1 remains highest on NQ (0.4382 vs. TRIO-8B's 0.4339).
• Small backbone is competitive but weaker out-of-domain. TRIO-4B (Qwen3-4B-Instruct) averages 0.3988, beating Search-R1 on HotpotQA (0.4174 vs. 0.4116) and 2Wiki (0.4373 vs. 0.4301) but falling behind on NQ (0.3878), Musique (0.1547), and Bamboogle (0.4640, tied with Search-R1). The authors attribute this to limited model capacity constraining generalization.
• Answer-agent training matters most. In the ablation, replacing the trained answer agent with the base backbone causes the largest drops: 14.14% (HotpotQA EM), 41.02% (2Wiki EM), 43.40% (Musique EM), and 43.55% (Bamboogle EM).
• Judge training helps least, and may overfit. TRIO w/o Judge Opt. is the strongest ablation variant and even gains 1.17% EM on 2Wiki over full TRIO, though it drops on Musique (14.22%) and Bamboogle (9.68%); the authors suggest the supervised judge training may overfit.
• Flexible stopping beats fixed retrieval rounds. On Musique, flexible rounds (max 3, judge-decided) reach 0.1857 EM / 0.2818 F1 versus the best fixed setting of 0.1585 EM / 0.2609 F1 (2 rounds). On Bamboogle, flexible gives 0.4960 EM / 0.6200 F1 versus 0.4560 / 0.5818 at 5 fixed rounds.
• Two-stage reward generalizes better. Training the answer agent with a single reward (R_a2) instead of the two-stage design yields a -5.9% average drop on two in-domain datasets and a -17.03% average drop on two out-of-domain datasets, despite higher training reward — read as overfitting to the training distribution.
• Efficient at inference. TRIO-8B averages 79.9 s per 100 queries versus 90.1 s for Search-R1, 88.0 s for R1-Searcher, 631.0 s for SuRe, and 695.8 s for DRAG. TRIO-4B averages 45.6 s and is described as achieving 98% of the performance of reasoning-based methods at roughly half the inference time.

Methodology in Plain English

The pipeline runs as a loop. An initial query retrieves chunks from a corpus. Then the query extension agent looks at the question, the queries issued so far, and the documents retrieved, and writes one new keyword-only query (maximum 8 keywords, no stopwords, must differ from existing queries). The judge agent then reads the accumulated queries and documents and outputs a binary decision: "Sufficient Information" or "Insufficient Information." If insufficient, the loop repeats up to a fixed budget; if sufficient, the answer agent produces the final answer from the current context only.

Each agent is trained separately, all supervision ultimately derived from question–answer pairs:

• Answer agent: Training samples come from a "retrieval-only debate" pipeline run up to three rounds on medium and hard multi-hop questions; snapshots of query pools and retrieved chunks at different sizes become training samples so the model learns to filter noise. It is trained with GRPO in two stages — first using the top 3 chunks per query with a reward of exact match plus a precision score plus a length-format penalty; then using the top 5 chunks per query with exact match plus token-level F1 plus the same format penalty.
• Query extension agent: The quality of a single query is defined as the change in the answer agent's F1 after adding the newly retrieved chunk versus before, and this difference is the GRPO reward. So the query agent is rewarded for retrieving evidence that actually improves answers.
• Judge agent: Trained by supervised fine-tuning on labels generated at test time — the answer agent produces 16 stochastic answers per sample; all 16 correct means "Sufficient," all 16 incorrect means "Insufficient," and mixed cases are discarded.

Why This Matters

The paper targets a practical weakness of RAG — that adding retrieval can make answers worse when queries are poorly formed or context is noisy — and shows that separating concerns across trained agents beats tuning a monolithic pipeline.

Impact on research: It provides a template for per-agent, task-specific optimization inside multi-agent RAG, moving beyond the shared-backbone assumption common in prior multi-agent RAG work, and offers a reward design (answer-quality delta) for evaluating individual retrieval queries.

Real-world applications:

• Domain question answering over private or specialized corpora such as healthcare and law, which the paper names as motivating settings where pre-trained models lack current knowledge.
• Multi-hop question answering over large document collections such as Wikipedia-scale corpora.
• Cost-sensitive deployments where inference budget matters, given TRIO's shorter output generation and lower per-query time than debate-style baselines.
• Enterprise search assistants that need to decide when to stop retrieving rather than always running a fixed number of rounds.

Industry relevance: The efficiency numbers (79.9 s/100 queries for TRIO-8B on two L40S GPUs) and the reuse of open backbones (Llama3.1-8B-Instruct, Qwen3-4B-Instruct) with an E5-base-v2 retriever make the approach practical to deploy and reproduce; code is released at https://github.com/yingtaoren/TRIO.

Future Directions

• Long-form and open-ended generation. The authors state TRIO is evaluated primarily on short-answer QA and has not been validated for long-form generation, where accumulated errors and irrelevant evidence across rounds could hurt.
• Live web search. Dynamic document availability and heterogeneous webpage structures are named as untested challenges.
• Better retrieval-stopping calibration. The limitations section calls for calibrating stopping decisions and evaluating failure modes in dynamic corpora; the ablation's slight regression on 2Wiki without a trained judge also hints at overfitting in the judge's supervised training.
• Closing the out-of-domain gap for smaller backbones. TRIO-4B's degradation on NQ, Musique, and Bamboogle raises the question of how far specialization can compensate for limited model capacity.

Target Audience

Researchers and engineers working on RAG, LLM agents, and reinforcement-learning-based post-training of language models. It suits readers who already understand standard retrieve-then-generate pipelines and want a concrete design for decomposing RAG into trained, cooperating roles; the ablation and retrieval-round analyses are also useful for practitioners deciding whether to invest in a trained query generator or a trained stopping judge.

Authors’ 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.