Context Tuning for Retrieval Augmented Generation
Authors: name: Raviteja Anantha; name: Tharun Bethi; name: Danil Vodianik; name: Srinivas Chappidi
Published: 2023-12-09 23:33:16

Abstract: Large language models (LLMs) have the remarkable ability to solve new tasks
with just a few examples, but they need access to the right tools. Retrieval
Augmented Generation (RAG) addresses this problem by retrieving a list of
relevant tools for a given task. However, RAG's tool retrieval step requires
all the required information to be explicitly present in the query. This is a
limitation, as semantic search, the widely adopted tool retrieval method, can
fail when the query is incomplete or lacks context. To address this limitation,
we propose Context Tuning for RAG, which employs a smart context retrieval
system to fetch relevant information that improves both tool retrieval and plan
generation. Our lightweight context retrieval model uses numerical,
categorical, and habitual usage signals to retrieve and rank context items. Our
empirical results demonstrate that context tuning significantly enhances
semantic search, achieving a 3.5-fold and 1.5-fold improvement in Recall@K for
context retrieval and tool retrieval tasks respectively, and resulting in an
11.6% increase in LLM-based planner accuracy. Additionally, we show that our
proposed lightweight model using Reciprocal Rank Fusion (RRF) with LambdaMART
outperforms GPT-4 based retrieval. Moreover, we observe context augmentation at
plan generation, even after tool retrieval, reduces hallucination.