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

TRIO. Task Refinement with Individual Optimization for Retrieval-Augmented Generation. By Yingtao Ren and five others, from the University of Technology Sydney, the University of Wisconsin-Madison, and one other institution. Published on the second of November.

Retrieval-augmented generation draws on external knowledge to improve what a language model produces. Progress has been considerable, yet current approaches remain far from satisfactory for two stated reasons. Retrieval is inadequate, and the model may still hallucinate. Either one can spoil an answer: the wrong evidence comes back, or the right evidence is not used well. A fix aimed at only one of those leaves the other in place. Better retrieval does not help if the answer still drifts from the evidence, and careful answering does not help if the evidence was never found.

TRIO, short for Task Refinement with Individual Optimization, splits the work among three agents. Each is well designed for one part of the task, and the three are placed inside an iterative refinement paradigm. The key is in the name. Every agent has its own individual optimization objective, and those objectives are aligned with one another so that improving each part improves the whole.

The first agent handles queries. Ambiguous queries are expanded by a query extension agent, and a reward is computed from how much the expansion improves the quality of the answer. That reward is used to refine the agent toward more efficient expansion.

The second agent is a judge, aimed at retrieval quality. A unanimity vote annotates each expanded query as either sufficient or insufficient, and the generative outcomes of that vote serve as ground truth for optimizing the judge. The third is an answer agent, responsible for using the evidence well and answering accurately. It is optimized with dynamic, matching-based rewards, combined with curriculum learning, to produce the final outputs.

Extensive experiments across multiple benchmark datasets compare TRIO with state-of-the-art approaches, and validate its effectiveness. The abstract does not report specific figures, so the size of the improvement, and which of the three agents contributes most to it, are not something the summary lets a listener judge.

Each of the three agents maps onto one of the weaknesses named at the outset. Query extension addresses ambiguity before retrieval begins, the judge addresses whether what was retrieved is enough, and the answer agent addresses how the evidence is turned into a reply. Giving each its own objective, while keeping those objectives aligned, is what lets the three be trained separately and still pull in one direction.