RoleRAG: Unified Retrieval-Augmented Generation

• RoleRAG is a unified retrieval-augmented generation framework that leverages role-specific soft prompts to efficiently manage diverse QA tasks within a single LLM.
• It integrates six modular components to decompose user queries into dynamic graphs, enhancing both accuracy and computational efficiency across various benchmarks.
• Parameter efficiency is achieved by training only the soft prompt embeddings while keeping the backbone model frozen, enabling rapid deployment and scalability.

RoleRAG is a unified retrieval-augmented generation (RAG) framework for handling multiple role-specific tasks within a single LLM instance, achieved via role-specific token optimization. It operationalizes a modular pipeline covering all key RAG sub-tasks, allows dynamic query decomposition using a directed acyclic query graph, and supports extensible, efficient multi-role reasoning and deployment. RoleRAG introduces a parameter-efficient approach by optimizing only soft prompt embeddings (role tokens), leaving the backbone model frozen, enabling dynamic module activation, and attaining strong empirical gains across a variety of open-domain question-answering benchmarks.

1. Framework Overview and Objectives

RoleRAG addresses the challenge of integrating diverse RAG sub-task optimizations—such as query decomposition, retrieval judgment, answer synthesis, and context distillation—into a single, efficient, deployable system. Distinct from conventional multi-stage or monolithic RAG pipelines, RoleRAG utilizes six specialized modules, each serving a granular RAG function, with all modules invoked via a single underlying LLM. This centralizes all RAG sub-task logic into a unified architecture, eliminating redundancy and streamlining scaling and deployment.

The system is motivated by two goals:

• Efficient multi-role fulfillment: Each RAG sub-task is framed as a distinct “role,” activated by its own learnable token prompt.
• Parameter efficiency and extensibility: By optimizing only per-role soft prompt embeddings (role tokens), and not fine-tuning the backbone, new tasks/modules can be deployed rapidly without large model retraining.

2. Modular Pipeline Components and Interactions

RoleRAG consists of six coordinated modules:

Module Name	Function	Activation Mechanism
Query Graph Builder	Decomposes user query into sub-queries/DAG nodes	Role tokens (soft prompt)
Retrieval Judge	Determines if retrieval is required for sub-query	Role tokens
Sub-answer Generator	Generates targeted answer for sub-query	Role tokens
Summarizer	Compresses retrieved context for downstream use	Role tokens
New Query Generator	Adds new sub-queries based on answer memory	Role tokens
Answer Reasoner	Synthesizes final, comprehensive answer	Role tokens

Interaction is orchestrated via a query graph (see Section 3), and answer memory (a persistent structure containing sub-query answers and summaries). The modules are not hard-wired into a single sequential pipeline; instead, the query graph structure and dynamic answer memory enable iterative resolution, adaptive pruning, or growth of sub-queries as new knowledge is synthesized or gaps are detected.

3. Query Graph Mechanism: Dynamic Decomposition and Execution

The query graph $G(Q)$ is a directed acyclic graph representing the decomposition of an initial natural-language user query $Q$ into a set of manageable sub-queries $\{q_1, ..., q_n\}$, each with parent/child relationships indicating execution dependencies.

• Construction: The Query Graph Builder (activated by its role token) decomposes $Q$ using LLM-driven analysis and outputs nodes, edges (dependencies), and possible placeholders (e.g., referencing previous answers in children nodes).
• Resolution: Sub-queries are executed via the Sub-answer Generator, with retrieval invoked as decided by the Retrieval Judge. Summaries of external context are built via the Summarizer.
• Dynamic Expansion: The New Query Generator may add nodes if the answer memory suggests knowledge gaps remain or higher-level reasoning steps are needed.
• Final Synthesis: The Answer Reasoner aggregates memory and resolved sub-queries to generate the final answer.

This structure supports robust, multi-hop reasoning, dynamic adjustment of search depth/width, and clear prevention of error compounding that occurs in purely iterative or monolithic pipelines.

4. Role-Specific Token Optimization

A core technical method is the use of trainable, module-specific soft prompt embeddings (role tokens). For each module (task/role), a unique segment of soft tokens is prepended to the input. Only these embeddings are trainable; the backbone LLM parameters ($\theta$) remain fixed. This is formalized as follows:

$$p = \prod_{i=1}^{m} p_{\theta, \delta}(y^T_i \mid X^T;\, t_1;\ldots;t_n;\, y^T_{<i})$$

Where:

• $X^T$ is the task input,
• $Y^T$ the output,
• $[t_1;...;t_n]$ are the trainable role tokens for task $T$,
• $\delta \in \mathbb{R}^{n \times d}$ are the token embeddings,
• $\theta$ (frozen backbone) is not updated.

This enables the following:

• All modules/roles share the same LLM weights, maximizing parameter reuse.
• Deployment only requires the frozen LLM plus a compact matrix of role tokens.
• Extending with new roles/modules is efficient—add new tokens and train only on relevant data.

Parameter efficiency is notable. E.g., for Llama-3-8B, 30 role tokens (4096 dimensions) require only ~0.1M parameters, a negligible fraction of the backbone size.

5. Experimental Results and Comparative Performance

RoleRAG achieves strong results across five open-domain QA datasets:

Dataset	Best Baseline F1	RoleRAG F1
HotpotQA	45.56 (BlendFilter)	49.17
MuSiQue	16.04 (RQ-RAG)	27.30
2WikiMultiHopQA	37.64 (RQ-RAG)	53.87
Bamboogle	41.10 (IRCoT)	54.47
PopQA	44.94 (SuRe)	45.42

• On multi-hop tasks (MuSiQue, 2WikiMultiHopQA), the explicit query decomposition and modularization yield improvements of up to 64% over leading baselines.
• For single-hop QA (PopQA), RoleRAG avoids overhead by constructing trivial graphs as needed.
• Ablations show removing the Query Graph Builder or Summarizer sharply degrades accuracy and efficiency; reducing retrieval (via Retrieval Judge) yields computational savings with minimal lost accuracy.

RoleRAG also shows robust out-of-domain generalization, outperforming baselines on datasets with compositional or adversarial queries.

6. Technical and Schematic Comparisons

Feature	Standard RAG	Iterative RAG	RoleRAG
Decomposition	Flat/simple	Monolithic	Query graph (modular, explicit)
Module specialization	None/partial	Joint/self-reflect	Per-module role tokens
Multi-role handling	Multiple LLMs	Single LLM	Single LLM with soft prompts
Parameter efficiency	N/A	Weak	Highest
Adaptivity	None	Limited	Dynamic via query graph
Out-of-domain robustness	Limited	Often weak	Strong
F1 (multi-hop QA)	~15–45	~16–33	49–54

7. Implications, Extensions, and Positioning

RoleRAG demonstrates that highly modular, explicit decomposition with dynamic query graph orchestration can be implemented in a resource-efficient manner using prompt-based, role-specialized control over a single backbone LLM. This architecture:

• Enables fine-grained resource and response management (retrieval as needed, summary compaction, new reasoning steps).
• Scales robustly to complex, multi-step reasoning settings without model duplication.
• Facilitates extensibility—new modules/roles require only prompt token additions.
• Achieves strong empirical gains on both in-domain and out-of-domain benchmarks.

A plausible implication is that such role-driven, soft-prompt LLM architectures may become foundational for highly modular, domain-adaptive LLM applications, enabling rapid iteration and robust performance even as tasks increase in diversity and complexity.