Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Single LLM, Multiple Roles: A Unified Retrieval-Augmented Generation Framework Using Role-Specific Token Optimization

Yutao Zhu, Jiajie Jin, Hongjin Qian, Zheng Liu, Zhicheng Dou, Ji-Rong Wen

TL;DR

RoleRAG presents a unified retrieval-augmented generation framework that uses role-specific token optimization to coordinate six modules within a single frozen LLM. By modeling query decomposition as a dynamic query graph and gating sub-tasks with a retrieval judge, it achieves efficient, multi-step reasoning and memory-driven synthesis for open-domain QA. Empirical results across five datasets show substantial gains over both sequential and iterative baselines, with strong generalization across model sizes and query complexity. The approach offers practical benefits in deployment efficiency and provides data-collection and training-objective strategies to support scalable multi-task prompting.

Abstract

Existing studies have optimized retrieval-augmented generation (RAG) across various sub-tasks, such as query understanding and retrieval refinement, but integrating these optimizations into a unified framework remains challenging. To tackle this problem, this work proposes RoleRAG, a unified RAG framework that achieves efficient multi-task processing through role-specific token optimization. RoleRAG comprises six modules, each handling a specific sub-task within the RAG process. Additionally, we introduce a query graph to represent the decomposition of the query, which can be dynamically resolved according to the decomposing state. All modules are driven by the same underlying LLM, distinguished by task-specific role tokens that are individually optimized. This design allows RoleRAG to dynamically activate different modules within a single LLM instance, thereby streamlining deployment and reducing resource consumption. Experimental results on five open-domain question-answering datasets demonstrate the effectiveness, generalizability, and flexibility of our framework.

Single LLM, Multiple Roles: A Unified Retrieval-Augmented Generation Framework Using Role-Specific Token Optimization

TL;DR

RoleRAG presents a unified retrieval-augmented generation framework that uses role-specific token optimization to coordinate six modules within a single frozen LLM. By modeling query decomposition as a dynamic query graph and gating sub-tasks with a retrieval judge, it achieves efficient, multi-step reasoning and memory-driven synthesis for open-domain QA. Empirical results across five datasets show substantial gains over both sequential and iterative baselines, with strong generalization across model sizes and query complexity. The approach offers practical benefits in deployment efficiency and provides data-collection and training-objective strategies to support scalable multi-task prompting.

Abstract

Existing studies have optimized retrieval-augmented generation (RAG) across various sub-tasks, such as query understanding and retrieval refinement, but integrating these optimizations into a unified framework remains challenging. To tackle this problem, this work proposes RoleRAG, a unified RAG framework that achieves efficient multi-task processing through role-specific token optimization. RoleRAG comprises six modules, each handling a specific sub-task within the RAG process. Additionally, we introduce a query graph to represent the decomposition of the query, which can be dynamically resolved according to the decomposing state. All modules are driven by the same underlying LLM, distinguished by task-specific role tokens that are individually optimized. This design allows RoleRAG to dynamically activate different modules within a single LLM instance, thereby streamlining deployment and reducing resource consumption. Experimental results on five open-domain question-answering datasets demonstrate the effectiveness, generalizability, and flexibility of our framework.

Paper Structure

This paper contains 24 sections, 1 equation, 13 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Module Design
  5. Data Collection
  6. Training Strategy
  7. Experiments
  8. Datasets and Evaluation Metrics
  9. Baselines
  10. Experimental Results
  11. Further Analysis
  12. Ablation Study
  13. Impact of Model Size
  14. Impact of Task-specific Token Amounts
  15. Impact of Query Complexity
  16. ...and 9 more sections

Figures (13)

  • Figure 1: Comparison between existing studies and our framework.
  • Figure 2: The illustration of our RoleRAG framework, which contains three main steps: (1) query graph building, (2) sub-query execution, and (3) final answer reasoning.
  • Figure 3: Illustration of the training and inference processes in our framework
  • Figure 4: Performance with various numbers of tokens.
  • Figure 5: Performance of various models on questions of different complexity (MuSiQue). "Ratio" indicates the proportion of a certain category to the entire data.
  • ...and 8 more figures