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

2025/05/21 by Yutao Zhu, Jiajie Jin, Zhu, Yutao +9
Computer Science · #Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #FOS: Computer and information sciences #Natural Language Processing Techniques #Semantic Web and Ontologies #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2505.15444

openalex publication_date 2025/05/21 · openalex created_date 2025/10/19 · openalex updated_date 2026/07/28

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.

Citations

• AutoRAG: Automated Framework for optimization of Retrieval Augmented Generation Pipeline
• The Llama 3 Herd of Models
• MindSearch: Mimicking Human Minds Elicits Deep AI Searcher
• TRACE the Evidence: Constructing Knowledge-Grounded Reasoning Chains for Retrieval-Augmented Generation
• RaFe: Ranking Feedback Improves Query Rewriting for RAG
• RQ-RAG: Learning to Refine Queries for Retrieval Augmented Generation
• Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity
• Small Models, Big Insights: Leveraging Slim Proxy Models To Decide When and What to Retrieve for LLMs
• BIDER: Bridging Knowledge Inconsistency for Efficient Retrieval-Augmented LLMs via Key Supporting Evidence
• BlendFilter: Advancing Retrieval-Augmented Large Language Models via Query Generation Blending and Knowledge Filtering
• OVM, Outcome-supervised Value Models for Planning in Mathematical Reasoning
• Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection
• Making Retrieval-Augmented Language Models Robust to Irrelevant Context
• Enhancing Retrieval-Augmented Large Language Models with Iterative Retrieval-Generation Synergy
• REPLUG: Retrieval-Augmented Black-Box Language Models
• Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions
• When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories
• Text Embeddings by Weakly-Supervised Contrastive Pre-training
• Solving math word problems with process- and outcome-based feedback
• Measuring and Narrowing the Compositionality Gap in Language Models
• Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps
• KILT: a Benchmark for Knowledge Intensive Language Tasks
• Language Models are Few-Shot Learners
• Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering
• KBM: Delineating Knowledge Boundary for Adaptive Retrieval in Large Language Models

Cited by

• Hierarchical Document Refinement for Long-context Retrieval-augmented Generation
• WebThinker: Empowering Large Reasoning Models with Deep Research Capability

Related

• Retrieval Augmented Generation (RAG) and Beyond: A Comprehensive Survey on How to Make your LLMs use External Data More Wisely
• Adapting Large Language Models for Multi-Domain Retrieval-Augmented-Generation
• Leveraging LLM-Assisted Query Understanding for Live Retrieval-Augmented Generation
• RAGO: Systematic Performance Optimization for Retrieval-Augmented Generation Serving
• Enhancing Retrieval and Managing Retrieval: A Four-Module Synergy for Improved Quality and Efficiency in RAG Systems