[論文レビュー] Reinforcement Learning for Optimizing RAG for Domain Chatbots

Mandar Kulkarni, Praveen Tangarajan+2 more|arXiv (Cornell University)|Jan 10, 2024

Topic Modeling被引用数 18

ひとこと要約

この論文はRAGパイプラインの外部にポリシー型RLアプローチを開発し、FAQ文脈の取得を判断します。GPT-4評価で約31%のトークン削減を達成し、ドメインFAQチャットボットの精度を若干向上させました。また、社内埋め込みモデルは取得とOOD検出で公開モデルを上回ることを示しています。

ABSTRACT

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.

研究の動機と目的

• RAG設定でのLLMトークンコストを削減することにより、効率的なドメイン固有のチャットボットを動機付ける。
• 情報NCEを用いた社内埋め込みモデルが、ドメインFAQの取得とOOD検出で公開埋め込みを上回ることを示す。
• 取得すべきFAQ文脈を決定するポリシーグラデーションRLエージェントを提案・評価し、コストを最小化する。
• RLベースの文脈選択と類似度閾値の組み合わせが、精度を損なうことなく大幅なトークン削減をもたらすことを示す。

提案手法

• ドメインFAQ取得のためのInfoNCE損失を用いた社内埋め込みモデルを訓練する。
• 英語およびヒンギリッシュのクエリに対するトップ1/トップ3取得精度で、社内埋め込みと公開モデルを比較する。
• 報酬評価者としてGPT-4を用い、Good/Bad評価をポリシー勾配訓練の数値報酬へ変換する。
• RAGとは外部の状態（前のクエリ、アクション、現在のクエリ）に基づいてFETCHまたはNO_FETCHアクションを選択するポリシーネットワークを開発する。
• 状態・アクション・報酬の軌跡を用いてポリシー勾配とエントロピー正則化でポリシーを訓練する。
• RLポリシーを類似度閾値（SimThr）と組み合わせて、トークン使用をさらに削減する。

Figure 1: Proposed policy agent based architecture for optimizing RAG for domain chatbots.

実験結果

リサーチクエスチョン

• RQ1 外部ポリモデルがFAQ文脈を取得するタイミングを学習し、回答品質を損なうことなくLLMトークン使用を削減できるか。
• RQ2 社内ドメイン適合埋め込みモデルは公開埋め込みより取得精度とOOD検出を改善するか。
• RQ3 RLベースの文脈選択は類似度閾値ルールとどのように相互作用してRAGコストを最適化するか。
• RQ4 RAG設定でポリシーを訓練するための自動評価者としてGPT-4を用いることの影響は何か。

主な発見

モデル	英語 top-1	英語 top-3	ヒンギリッシュ top-1	ヒンギリッシュ top-3
e5-base-v2	0.82	0.91	0.71	0.87
triplet-loss	0.90	0.93	0.84	0.89
infoNCE	0.97	0.98	0.94	0.95

• InfoNCEで学習した社内埋め込みモデルは、英語とヒンギリッシュのクエリに対して公開のe5-base-v2よりトップ1/トップ3の精度が高い。
• 社内モデルはドメイン内およびOOD識別で優れており、類似度閾値を用いて取得を選択的にスキップできる。
• RAG外部のRLポリシーは、類似度閾値と組み合わせることで、91クエリのテストセッションにおいてトークン使用を約31%削減し、精度がわずかに向上した。
• GPT-4の評価は報酬へ変換してポリシーグラデ更新を推進できる。
• GPT-2をポリシーモデルとして使用してもトークン削減が得られ（約25%）、このアプローチはポリシーアーキテクチャを超えて一般化可能である。
• 異なる報酬設計はトークン削減に影響を与える可能性がある（例：別の設計で約30%）。