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

TL;DR

This paper introduces RoleRAG, a unified retrieval-augmented generation framework that uses role-specific token optimization within a single LLM to efficiently handle multiple sub-tasks in open-domain question answering.

Contribution

The paper presents RoleRAG, a novel unified framework that integrates multiple RAG sub-tasks into one model using role tokens, enhancing efficiency and flexibility.

Findings

Effective on five open-domain QA datasets

Reduces resource consumption compared to separate modules

Demonstrates high generalizability and flexibility

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.