VRSD: Rethinking Similarity and Diversity for Retrieval in Large Language Models

Retrieval methods in Large Language Models (LLMs) have gained traction, particularly due to their pivotal role in open-domain question answering, evidenced by seminal contributions in the field [5, 9, 14, 25, 10]. The introduction of Retrieval Augmented Generation (RAG) further underscored the significance of these methods across knowledge-intensive tasks [16], notably enhancing generation capabilities for open-domain queries [20]. This spurred the adoption of techniques such as K-Nearest-Neighbor (KNN) in diverse applications, ranging from the customization of multilingual models in machine translation [12] to improving the prediction of rare patterns in LLMs [13, 1]. Continued advancements in retrieval techniques have focused on identifying highly informative examples to augment in-context learning, thereby enabling LLM-based systems to achieve significant performance improvements with minimal examples [3]. Early on, traditional sparse retrieval methods like BM25 [28]—an extension of TF-IDF—were utilized to refine in-context learning [17]. Subsequently, the integration of LLMs’ intrinsic capabilities [30] and Sentence-BERT (SBERT) [26] facilitated the retrieval of highly pertinent examples for prompt integration. The advent of dense retrievers signified a methodological enhancement in retrieval from a machine learning perspective, and the incorporation of feedback signals with contrastive learning has yielded more effective retrieval systems [29, 32]. Recent innovations like UPRISE [6] and PRAC [23] have further optimized the performance of in-context learning by retrieving demonstrations directly from training data. Despite these advances, most retrieval methods still treat each candidate independently, which can lead to suboptimal outcomes due to the interaction effects among in-context examples, resulting in a lack of diversity. Given the expanding applications of LLMs, diversity becomes increasingly crucial, as complex inference tasks often require undefined example types. Incorporating a diverse range of examples enriches LLMs’ learning processes, facilitating more innovative and robust responses, especially for complex open-ended questions. Furthermore, a relevance-focused diversity retrieval method could help mitigate the impact of malicious information as LLMs continue to scale and assimilate more societal data. Even if some malicious information is inadvertently retrieved, the inclusion of diverse data within the same batch can enable LLMs to fully comprehend the query’s intent and provide accurate responses. Additionally, as LLMs are tasked with processing ever-greater volumes of information, the necessity for a straightforward and efficient retrieval method intensifies, minimizing the dependency on resource-intensive preparatory tasks.

To enhance retrieval processes by accounting for both relevance and diversity, the Maximal Marginal Relevance (MMR) algorithm was introduced [4]. MMR addresses the balance between relevance and diversity in traditional retrieval and summarization methods by employing "marginal relevance" as an evaluation metric. This metric is defined as a linear combination of independently measured relevance and novelty, formulated as Eq.1:

The challenge lies in selecting an appropriate $\lambda$ to achieve the desired balance between relevance and diversity, particularly in high-dimensional vector spaces where the impact of varying $\lambda$ is less predictable. This variability in $\lambda$ leads to fluctuations in retrieval results, resulting in unpredictable consequences, which can be illustrated by a simple example. Commonly, $\lambda$ is preset at a value of 0.5 in many MMR implementations, a choice that stems from the algorithm’s foundational design. It is important to note that at $\lambda=1$, the algorithm exclusively prioritizes relevance, while at $\lambda=0$, it focuses entirely on diversity. Let us examines the performance of the MMR algorithm at the typical midpoint setting of $\lambda=0.5$. For clarity and ease of comprehension, we model the retrieval process within a two-dimensional vector space, though the principles observed are equally applicable to more complex, higher-dimensional scenarios.

As illustrated in Figure.1(a), consider $q$ as the query vector, and $d_{0}$ to $d_{3}$ as candidate vectors that surpass the relevance threshold, collectively represented as $R=\{d_{0},d_{1},d_{2},d_{3}\}$, with $S$ initially empty. Utilizing the MMR algorithm, $d_{0}$ is first selected due to its highest relevance to $q$, determined using cosine similarity as a measure. Subsequently, $d_{3}$ is chosen over $d_{1}$, despite $d_{1}$ having a smaller angle with $q$ and thus greater direct relevance. The selection of $d_{3}$ is influenced by the fact that the cumulative relevance between $d_{1}$ and $d_{0}$ significantly surpasses that between $d_{3}$ and $d_{0}$, resulting in a higher MMR value for $d_{3}$ as per the formula.

However, as depicted in Figure.1(b), with $q$ serving as the query vector and $R=\{d_{0},d_{1},d_{2},d_{3}\}$ representing the initial set of candidate vectors, $d_{0}$ is first selected due to its maximal relevance to $q$. The selection process using the MMR algorithm proceeds as follows: with $\lambda=0.5$, $S=\{d_{0}\}$, and $R\setminus S=\{d_{1},d_{2},d_{3}\}$, the formula can be articulated as Eq.2:

Given that $d_{0}$, $d_{1}$, $d_{2}$, and $d_{3}$ are positioned on the same side relative to $q$, and assuming both $\text{Sim}_{1}$ and $\text{Sim}_{2}$ denote cosine similarity, let $\theta$ represent the angle between $d_{0}$ and $q$, and $x$ denote the angle between $d_{i}$ (i.e., $d_{1},d_{2},d_{3}$) and $d_{0}$. Thus, we get the Eq.3

The function $f(x)=\cos(x+\theta)-\cos(x)$, with its derivative $f^{\prime}(x)=-\sin(x+\theta)+\sin(x)$, assumes $x$ and $x+\theta$ lie within $(0,\pi/2)$. Consequently, $f^{\prime}(x)<0$, indicating that for vectors on the same side of $q$, their MMR values decrease as the angle with $q$ increases. Thus, following the selection of $d_{0}$, the subsequent choices are $d_{1}$, then $d_{2}$, and so on. This sequence suggests that relevance predominantly influences the selection outcome.

The real challenge in vector retrieval emerges when $\lambda\neq 0.5$. The selection among candidate vectors $d_{1}$, $d_{2}$, and $d_{3}$ hinges critically on both $\lambda$ and $\theta$, complicating the determination of the most appropriate candidate. This dependency means that different query vectors and the distribution of initial candidate vectors require varying $\lambda$ values to achieve optimal performance. Consequently, it is impractical to predict the value of $\lambda$ in advance or to ascertain a precise direction for optimization. This issue becomes even more pronounced in higher-dimensional vector spaces, where the perturbations induced by changing $\lambda$ complicate the identification of an optimal adjustment direction. This inherent complexity underscores the need for adaptive retrieval strategies that dynamically adjust $\lambda$ based on the characteristics of the query and candidate vector distributions.

To address the problem of selecting a subset of vectors from a set of candidate vectors that satisfy both similarity and diversity requirements, we refer to the MMR algorithm and several LLM-based algorithms, typically considering the following premises: Firstly, the candidate vectors are identified from the entire set of vectors ($size=N$) using similarity metrics, resulting in a subset of vectors ($size=n$). Consequently, this set of candidate vectors inherently exhibits a relative high degree of similarity to the query vector. Secondly, within these $n$ candidate vectors, the vector most similar to the query vector is typically selected first, as is the case with the MMR algorithm and others. This approach is favored because, in applications such as in-context learning with LLMs, examples with the highest similarity to the query are generally the most helpful.

As previously mentioned, while algorithms like MMR are widely applied in practice, these studies often lack a robust and reliable theoretical model. In other words, many approaches employ heuristic strategies or machine learning methods to arrive at a solution without providing a rigorous formal description and analysis of similarity and diversity from a theoretical perspective. Therefore, based on the aforementioned premises, we propose using the sum vector to characterize both similarity and diversity in vector retrieval. The definition of the sum vector is as follows:

Specifically, we aim to maximize the similarity between the sum vector of the selected $k$ vectors and the query vector. On one hand, maximizing the similarity of the sum vector to the query vector imposes a similarity constraint. On the other hand, from a geometric perspective, ensuring the sum vector is similar to the query vector means that the selected vectors approach the query vector from different directions, thus imposing a diversity constraint. Therefore, using the sum vector to characterize similarity and diversity allows us to model complex semantic similarity and diversity through simple vector addition operations. Next, we define the problem of vectors retrieval as follows:

The vector $d_{0}$, characterized by its maximal similarity to the query vector $q$, establishes an initial constraint on similarity. The ensuing optimization objective strives to maximize the cosine similarity between the sum vector of all selected vectors and $q$. This process necessitates the selection of vectors that not only converge towards $q$ from diverse dimensions but also exhibit significant diversity and complementarity. However, upon further examination of above problem, we find that it is an NP-complete problem. Below, we provide a theoretical proof. Since the vector $d_{0}$, with the highest similarity, is initially selected, the subsequent selection of $k-1$ vectors must have the maximum cosine similarity with $q-d_{0}$. That is, maximizing the similarity between sum vector $d=d_{0}+d^{\prime}_{1}+d^{\prime}_{2}+...+d^{\prime}_{k-1}$ ($d^{\prime}_{1},d^{\prime}_{2},...,d^{\prime}_{k-1}$ represents the $k-1$ vectors selected subsequently) and $q$, is equivalent to maximizing the similarity between $d^{\prime}=d^{\prime}_{1}+d^{\prime}_{2}+...+d^{\prime}_{k-1}$ and $q-d_{0}$. To this end, we define a decision problem, namely:

Next, we will prove this decision problem is NP-complete. For the sake of concise proof, we further restrict the components of vectors to integers. The proof strategy is to reduce the subset sum problem [7] to this decision problem.

For the convenience of proof, we also need to define a modified version of the subset sum problem, called the $k$-subset sum problem.

Now it is time to prove the NP-completeness of the decision problem of vectors retrieval.

Since the vectors retrieval problem $(R,q,k)$ is a NP-complete problem, necessitating the use of heuristic methods to derive feasible solutions. Specifically, given a set of candidate vectors with high similarity, the objective is to select $k$ vectors that maximize the cosine similarity between the sum vector of the $k$ selected vectors and the query vector. We propose a new algorithm denoted as Vectors Retrieval with Similarity and Diversity (VRSD). VRSD initially selects the vector most similar to the query vector and then iteratively selects additional vectors from the remaining candidates. In each iteration, it chooses the vector that maximizes the cosine similarity between the cumulative sum of all previously selected vectors and the query vector, continuing this process until $k$ vectors are chosen. Further details about the VRSD algorithm can be found in Algorithm.1 and Fig.2.

As depicted in Algorithm.1, the time complexity of the VRSD algorithm is $k\times|R|=k\times n$, which accounts for the initial step of selecting $n$ candidate vectors from the entire set of vectors (size = $N$) based on similarity. Given that $N\gg n>k$, the computational load of subsequent steps in Algorithm.1 is minimal in comparison. The MMR algorithm, which also selects $k$ vectors from $|R|$ candidates, requires two iterations of maximum calculations as depicted in Eq.1—once for each candidate vector against the query vector and once against the set of already selected vectors $|S|$. Thus, the complexity for MMR becomes $k\times|R|\times|S|=k\times|R|^{2}=k\times n^{2}$, indicating a marginally higher computational demand compared to VRSD.

We evaluated the VRSD algorithm using three publicly available datasets of different categories and compared the VRSD with the MMR algorithm when the values of $\lambda$ are 0, 0.5, and 1 respectively :

• •

  ARC-DA [2]: A dataset of direct-answer science questions derived from the ARC multiple-choice question. Each example contains a question and multiple answers.

• •

  OpenBookQA [21]: A dataset of multiple-choice science questions, which probe the understanding of science facts and the application of these facts to novel situations. Each example contains a question, multiple choices, and an answer.

• •

  Puzzle [18]: A question answering dataset. These questions belong to lateral thinking puzzle. Each example contains a question and an answer.

For each item in each datasets, we concatenate the question part with its corresponding answer, subsequently selecting 20% of these concatenated items to form the test set, wherein the question parts are isolated. Items designated for the test set are excluded from the original dataset for subsequent experiments, where four examples are retrieved for each test question. Our evaluation focuses on two primary aspects: retrieval quality and answer quality. Retrieval quality is assessed by aggregating four vectors retrieved using either VRSD or MMR into a sum vector—denoted as $d_{\text{VRSD}}$ and $d_{\text{MMR}}$—which reflects the vectorial direction from which the examples approach the query vector $q$. We compute the cosine similarity between the sum vectors and the query vector as $\cos(d_{\text{VRSD}},q)$ and $\cos(d_{\text{MMR}},q)$. The comparison includes counting instances where $\cos(d_{\text{VRSD}},q)$ exceeds $\cos(d_{\text{MMR}},q)$, termed as the VRSD win. rate, and calculating the maximum difference between these cosine similarities for all queries in each test set. Additionally, we compute the mean cosine similarity values for these vectors. Such metrics are instrumental in elucidating the algorithms’ capacity to balance similarity and diversity. For answer quality assessment, we reconstruct prompts by concatenating the original sentences corresponding to the four retrieved vectors with the initial question, and input these into large LLMs. The responses generated by open-source LLM open-mistral-7b and closed-source LLM gpt-3.5-turbo are then compared with standard answers using the Rogue-L metric for ARC-DA and Puzzle datasets, and Accuracy for OpenBookQA. This evaluation helps us ascertain the efficacy of the retrieved examples in providing accurate answers

Table.1 presents the results of retrieval quality, indicating that the win rate of VRSD, compared to MMR across various datasets and conditions, consistently exceeds 90%. This outcome suggests that VRSD not only retrieves examples more relevant to the original query but also better satisfies the diversity requirements. Remarkably, VRSD maintains a minimum advantage of 0.05 over MMR in the worst-case scenario concerning the maximum difference. Additionally, the mean value of cosine similarity between the sum vector and the query vector is significantly higher for VRSD than for MMR. VRSD generally demonstrates superior retrieval efficacy from diverse perspectives. Table 2 shows the answer quality results, with VRSD achieving the highest scores across all metrics, thus indicating that the examples retrieved by VRSD enhance the LLM’s understanding of the query and facilitate the generation of the desired answers. The superior performance of the closed-source LLM over the open-source LLM can be attributed to the former’s enhanced understanding and reasoning capabilities. Overall, VRSD exhibits superior performance to MMR in both example retrieval and query answering, fulfilling the requirements of both similarity and diversity without the need for parameter adjustments, thereby highlighting the pronounced advantages of VRSD.

It is worth noting that, compared to other datasets, the mean values of $\cos(d_{\text{VRSD}},q)$ and $\cos(d_{\text{MMR}},q)$ in the Puzzle dataset are relatively low. Additionally, the difference in results between the open-source and closed-source models is less pronounced, and VRSD achieves a 100% win rate, which may be attributed to the small size of the dataset. Nevertheless, given that lateral-puzzle questions require LLMs to comprehend the query and derive insights from various angles in the retrieved examples, the Puzzle dataset remains a valuable tool for assessing our algorithm.