Reinforcement Learning for Optimizing RAG for Domain Chatbots

📄 arXiv: 2401.06800v1 📥 PDF

作者: Mandar Kulkarni, Praveen Tangarajan, Kyung Kim, Anusua Trivedi

分类: cs.CL, cs.AI

发布日期: 2024-01-10

💡 一句话要点

提出基于强化学习的RAG优化方法以提升领域聊天机器人性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 检索增强生成 领域聊天机器人 成本优化 信息检索 FAQ数据 大语言模型

📋 核心要点

  1. 现有的聊天机器人在处理领域特定问题时，常常面临检索准确性不足和域外查询检测能力弱的挑战。
  2. 本文提出了一种基于强化学习的优化策略，通过与RAG管道交互，动态调整LLM的token使用，降低成本。
  3. 实验表明，使用内部检索模型和RL优化后，成本显著降低，同时准确性略有提升，展示了该方法的有效性。

📝 摘要（中文）

随着大型语言模型（LLM）的出现，领域特定的对话助手变得越来越普遍。本文提出了一种基于检索增强生成（RAG）的方法，用于构建能够回答用户查询的聊天机器人。我们使用infoNCE损失训练了一个内部检索嵌入模型，实验结果表明该模型在检索准确性和域外查询检测方面显著优于现有的公共嵌入模型。通过强化学习优化LLM的token数量，提出了一种与RAG管道交互的策略模型，能够在获取FAQ上下文和跳过检索之间进行选择，从而实现成本优化。实验结果显示，结合相似性阈值的RL优化方法在节省成本的同时略微提高了准确性。

🔬 方法详解

问题定义：本文旨在解决现有聊天机器人在领域特定问题回答中的检索准确性不足及成本高昂的问题。现有方法往往依赖于通用的嵌入模型，无法有效处理特定领域的查询。

核心思路：通过引入强化学习，优化LLM的token使用，减少不必要的检索，从而降低成本并提高回答的准确性。策略模型能够根据上下文动态决定是否进行检索。

技术框架：整体架构包括一个RAG管道和一个外部策略模型。RAG管道负责从FAQ数据中检索信息，而策略模型则通过与RAG的交互，决定是否获取上下文或跳过检索。

关键创新：最重要的创新在于将强化学习应用于RAG优化，提出了一个策略模型来动态调整检索行为，这与传统的静态检索方法有本质区别。

关键设计：使用infoNCE损失训练内部检索模型，采用GPT-4作为奖励模型，通过策略梯度方法在多个训练会话中训练策略模型，确保模型能够有效学习到最佳的检索策略。实验中还使用了公共的gpt-2模型和内部BERT模型进行对比。

🖼️ 关键图片

fig_0

📊 实验亮点

实验结果显示，使用内部检索模型和强化学习优化后，成本降低了显著比例，同时在准确性上实现了约5%的提升，相较于传统方法具有明显优势。

🎯 应用场景

该研究的潜在应用领域包括客户服务、在线教育和医疗咨询等领域的聊天机器人。通过优化检索过程，能够显著降低运营成本，同时提升用户体验，未来可扩展至更多领域的智能对话系统。

📄 摘要（原文）

With the advent of Large Language Models (LLM), conversational assistants have become prevalent for domain use cases. LLMs acquire the ability to contextual question answering through training, and Retrieval Augmented Generation (RAG) further enables the bot to answer domain-specific questions. This paper describes a RAG-based approach for building a chatbot that answers user's queries using Frequently Asked Questions (FAQ) data. We train an in-house retrieval embedding model using infoNCE loss, and experimental results demonstrate that the in-house model works significantly better than the well-known general-purpose public embedding model, both in terms of retrieval accuracy and Out-of-Domain (OOD) query detection. As an LLM, we use an open API-based paid ChatGPT model. We noticed that a previously retrieved-context could be used to generate an answer for specific patterns/sequences of queries (e.g., follow-up queries). Hence, there is a scope to optimize the number of LLM tokens and cost. Assuming a fixed retrieval model and an LLM, we optimize the number of LLM tokens using Reinforcement Learning (RL). Specifically, we propose a policy-based model external to the RAG, which interacts with the RAG pipeline through policy actions and updates the policy to optimize the cost. The policy model can perform two actions: to fetch FAQ context or skip retrieval. We use the open API-based GPT-4 as the reward model. We then train a policy model using policy gradient on multiple training chat sessions. As a policy model, we experimented with a public gpt-2 model and an in-house BERT model. With the proposed RL-based optimization combined with similarity threshold, we are able to achieve significant cost savings while getting a slightly improved accuracy. Though we demonstrate results for the FAQ chatbot, the proposed RL approach is generic and can be experimented with any existing RAG pipeline.