SmartRAG: Jointly Learn RAG-Related Tasks From the Environment Feedback

TL;DR

SmartRAG introduces a jointly optimized pipeline for RAG systems, integrating modules like retrieval, query rewriting, and answer generation through reinforcement learning to enhance overall performance and efficiency.

Contribution

This work presents a novel joint optimization framework for RAG modules, improving their coordination and performance compared to traditional separately trained systems.

Findings

Jointly optimized SmartRAG outperforms separately trained systems.

Reinforcement learning effectively coordinates modules for better results.

System reduces retrieval costs while maintaining high accuracy.

Abstract

RAG systems consist of multiple modules to work together. However, these modules are usually separately trained. We argue that a system like RAG that incorporates multiple modules should be jointly optimized to achieve optimal performance. To demonstrate this, we design a specific pipeline called \textbf{SmartRAG} that includes a policy network and a retriever. The policy network can serve as 1) a decision maker that decides when to retrieve, 2) a query rewriter to generate a query most suited to the retriever, and 3) an answer generator that produces the final response with/without the observations. We then propose to jointly optimize the whole system using a reinforcement learning algorithm, with the reward designed to encourage the system to achieve the best performance with minimal retrieval cost. When jointly optimized, all the modules can be aware of how other modules are working and thus find the best way to work together as a complete system. Empirical results demonstrate that the jointly optimized SmartRAG can achieve better performance than separately optimized counterparts.