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

Abstract

Existing studies have optimized retrievalaugmented 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.