Better RAG using Relevant Information Gain

A limitation of transformer-based Large Language Models (LLMs) is that the number of tokens is bounded by the transformer’s context window, which is typically in the thousands. This is often insufficient for representing large texts, such as novels and corporate documentation. A common way to mitigate this constraint is via retrieval augmented generation (RAG), in which a relatively small subset of relevant passages are retrieved from a larger database and inserted into an LLM’s context window Gao et al. (2024). Typically, this process involves applying a similarity metric, such as cosine similarity, to (precomputed) embeddings of passages and the embedding of a query. Using this metric, many systems then use K-nearest-neighbors or a fast approximation with a vector database such as FAISS Douze et al. (2024). Importantly, K-nearest-neighbors Bijalwan et al. (2014) and related methods (such as a cross-encoder reranker Nogueira and Cho (2020)) simply return the highest individually relevant passages, without regard to whether the information in the passages is redundant. Given the premium value on LLM context-window real estate, it’s important to make best use of this limited resource by minimizing redundancy, while maintaining relevance.

To appreciate the importance of minimizing redundancy in a RAG context, consider a toy database of facts and the two possible sets of retrieval results in Table 1, for the same query, “Tell me some facts about sharks.” Both sets of retrieved results are highly relevant to the query, but only the second set is diverse enough to support a satisfactory answer.

A family of methods from the Information Retrieval literature attempts to address the general issue of diversity in retrieved results by introducing a measure that explicitly balances diversity and relevance Carbonell and Goldstein (1998). In this paper, we propose a more principled method, Dartboard, that instead seeks to directly accomplish what previous methods are indirectly aiming for - maximize the total amount of information relevant for a given query in a set of $k$ results. The intuition behind Dartboard is simple - we assume that one passage is the “correct” one for a given query. Our system is allowed $k$ “guesses” and it aims to maximize the relevance score of its most relevant guess. Since the best guess is not known ahead of time, this score is weighted by the probability of that guess being the most relevant. This objective is sufficient to encourage diversity in the guesses. This is because a redundant guess does little to increase the relevance of the most relevant guess.

The main contributions of this paper are 3-fold:

• •

  We introduce the Dartboard algorithm, a principled retrieval method based on optimizing a simple metric of total information gain relevant to a given query (§2).

• •

  We demonstrate the effectiveness of Dartboard on Retrieval-Augmented Generation Benchmark (RGB) Chen et al. (2023), a closed-domain question answering task. This benchmark consists of a retrieval component, and an end-to-end question-answering component. We show that the Dartboard algorithm, when used as the retrieval component, outperforms all existing baselines at both the component level and at end-to-end level (§3.1).

• •

  We show that instead of directly encouraging diversity, diversity naturally emerges by optimizing this metric (§A.5).

The Dartboard algorithm is based on the following analogy illustrated in Figure 1: Suppose that we have a cooperative two-player game where a dartboard is covered with a random collection of points. Player 1 is given one of these points arbitrarily as the target. Player 1 then throws her dart aiming for the target, and it lands somewhere on the board. Where it lands is the query. Player 2 sees where Player 1’s dart landed (the query), but doesn’t know where the actual target is. Player 2 then picks $k$ of the points on the board. The true target is revealed, and the score (which the players are trying to minimize) is the distance from the target to the closest guess. Note that to minimize the score, Player 2 would not want to put all his guesses right next to each other. Also, Player 2 should take into account how accurate Player 1’s throws are in general. In our implementation, Player 1’s accuracy is modeled by a Gaussian distribution with standard deviation $\sigma$.

More formally, Player 1 selects a target $T$ from a set of all points $A$ and gives a query $q$. Then Player 2 makes a set of guesses $G\subseteq A$, resulting in a score $s\left(G,q,A,\sigma\right)$ which is given as:

where $D$ is a distance function. For $d$ dimensional vectors, $A\subseteq\mathbb{R}^{d}$; under some assumptions, we can use a Gaussian kernel for the distance functions. For example, we can set $P\left(T=t|q,\sigma\right)=\mathcal{N}\left(q,t,\sigma\right)$. Thus, our equation becomes:

We tested Dartboard on benchmark datasets from Chen et al. (2023), from which we used two types of closed-domain question answering. In the simple question answering case, a query is answerable from a single passage retrieved from the corpus. For example, consider the query When is the premiere of ‘Carole King & James Taylor: Just Call Out My Name’?. On the other hand, in the information integration case, a query would require multiple passages to be retrieved to answer the query. For example, consider the query Who is the director of ‘Carole King & James Taylor: Just Call Out My Name’ and when is its premiere?. We modified this benchmark for our setup in the following way. The original benchmark contains “positive” and “negative” labeled passages for each query. The positive passages are useful for answering, while the negative ones are related but ineffective in answering the query. Since we are interested in the retrieval component of this task, we merged the positive and negative passages for all queries into a single collection of $11,641$ passages for the $300$ simple question answering test cases and $5,701$ passages for the $100$ information integration test cases. The evaluation is otherwise identical apart from the retrieval component. Note that the innovation of Dartboard is solely on the retrieval component. Therefore, we keep the rest of the RAG pipeline fixed. In particular, we do not modify the prompting of LLMs or try to optimize passage embeddings.

Given a query and the full set of thousands of passage embeddings, we measured both a direct retrieval score and the overall end-to-end performance of the system with the only change being the retrieval algorithm. For the direct retrieval score, we computed the Normalized Discounted Cumulative Gain (NDCG) score Wang et al. (2013) on retrieving any one of the “positive” passages relevant to a specific query. In the information integration case, the positive passages were split into positive ones for each component of the question. Therefore, in this case, we calculated the NDCG score for retrieving at least one positive passage for each component of the query. For the end-to-end score, given an LLM’s response to the query (generated from retrieved passages), we use the same evaluation as Chen et al. (2023), which does a string match of the response on a set of correct answers, marking each response as either correct or incorrect.

Some of the methods (described in Appendix A.2), including Dartboard, have tunable parameters. For instance, Maximal Marginal Relevance (MMR) has a diversity parameter that varies from 0 to 1. We performed a grid search over these parameters, reporting the best results for each method.

MMR retrieves documents Carbonell and Goldstein (1998) that are both relevant to the query and dissimilar to previously retrieved documents. It combines a relevance score (e.g., from BM25) with a novelty score that penalizes documents similar to those already retrieved. It have been used extensively for building recommendation systems Xia et al. (2015); Wu et al. (2023) as well as for summarization tasks Agarwal et al. (2022); Adams et al. (2022). However, MMR suffers from few limitations. First is that MMR requires the diversity parameter to control the balance between relevance and novelty. This parameter is often dataset-specific and requires careful tuning, making it impractical for real-world applications. Second is that MMR can favor exact duplicates of previously retrieved documents as they retain a high relevance score while minimally impacting the average novelty score (See Appendix A.7).

KNN retrieves documents based on their similarity to a query embedding Dharani and Aroquiaraj (2013); Bijalwan et al. (2014). While efficient, KNN often suffers from redundancy as nearby documents in the embedding space tend to be semantically similar Taunk et al. (2019). This can lead to a retrieved set dominated by passages conveying the same information with slight variations.

Several recent works have explored incorporating diversity objectives into retrieval models Angel and Koudas (2011); Li et al. (2015); Fromm et al. (2021). These approaches often involve complex optimization functions or require additional training data for diversity estimation. For example, Learning-to-Rank with Diversity methods leverage learning-to-rank frameworks that incorporate diversity objectives directly into the ranking function. This allows for the optimization of both relevance and diversity during the ranking process. However, these approaches often require large amounts of labeled training data for diversity, which can be expensive and time-consuming to obtain Wasilewski and Hurley (2016); Yan et al. (2021). Bandit-based approaches model document selection as a multi-armed bandit problem Hofmann et al. (2011); Wang et al. (2021). The model explores different retrieval strategies and receives feedback based on the relevance and diversity of the retrieved passages. These approaches can be effective but can be computationally expensive for large-scale retrieval tasks.

RAG models have also been extended to incorporate diversity objectives. For example, RAG with Dense Passage Retrieval retrieves a large number of candidate passages Cuconasu et al. (2024); Reichman and Heck (2024); Siriwardhana et al. (2023). It then employs a two-stage selection process: first selecting a diverse subset based on novelty scores, then selecting the most relevant passages from this subset. While effective, this approach requires careful tuning of the selection thresholds.

In this paper, we introduce Dartboard, a principled retrieval algorithm that implicitly encourages diversity of retrieved passages by optimizing for relevant information gain. We demonstrate that Dartboard outperforms existing state-of-the-art retrieval algorithms on both retrieval and end-to-end QA tasks. We view this work as an initial step for a more general line of work that optimizes information gain during retrieval, especially in the context of RAG systems. In future work, we plan to investigate Dartboard for other retrieval tasks, such as suggestion generation (see Appendix A.6).

We have not done a systematic investigation of the run time of Dartboard. In the worst case scenario, Dartboard is quadratic in the number of ranked passages. However, in practice, Dartboard hybrid typically runs in a fraction of a second for ranking (based on cosine-similarity with query) a set of $100$ passages (note that a full cross-encoder based MMR/Dartboard needs to run the cross-encoder $10,000$ times, and can take several seconds). This retrieval time is minimal compared to the time required for a LLM to process the retrieved passages and generate an answer.

Our experimental results are limited to a single benchmark and a single LLM i.e. ChatGLM Hou et al. (2024). It remains to be seen whether our results would generalize to other benchmarks and LLMs. We plan to investigate this in future work.

One shortcoming of our method (also shared by MMR) is that it requires a hyperparameter that affects how much diversity is encouraged. While we show that Dartboard is robust to the choice of this hyperparameter, it would be ideal to have a method that does not require manual tuning. As part of future work, we plan to investigate methods that automatically adapt to the context of the query. For example, the hyperparameter could be set based on a held-out validation set.

Another topic for future work is to investigate if it is also possible for $\sigma$ to vary depending on the type of query. For example, a query like “Tell me facts about The Beatles” would warrant a broader range of passages than a query like “Tell me facts about George Harrison”.

Another shortcoming of our approach is that our benchmarking criteria is limited in terms of the evaluation protocol we are using. Our evaluation is based on an exact string match of the output answer generated from the LLM with a set of possible answers. For example, for one question, the generated output answer is considered correct if it contains the exact string ‘January 2 2022’, ‘Jan 2, 2022’, etc., but would be considered incorrect if it only contains ‘January 2nd, 2022’. However, we left the benchmark as is (modulo our modifications mentioned above) so that our method is easily comparable to that of others.

Finally, though the initial cosine similarity based proposed Dartboard method is principled, the hybrid variation of Dartboard is not that principled. This is because it tries to compare logits from a cross-encoder with the cosine similarity of a different embedding model, similar to comparing apples with oranges, though it seems to work well as seen in our presented empirical results.