Question-Based Retrieval using Atomic Units for Enterprise RAG (2024)

Abstract: Enterprise retrieval augmented generation (RAG) offers a highly flexible framework for combining powerful large language models (LLMs) with internal, possibly temporally changing, documents. In RAG, documents are first chunked. Relevant chunks are then retrieved for a user query, which are passed as context to a synthesizer LLM to generate the query response. However, the retrieval step can limit performance, as incorrect chunks can lead the synthesizer LLM to generate a false response. This work applies a zero-shot adaptation of standard dense retrieval steps for more accurate chunk recall. Specifically, a chunk is first decomposed into atomic statements. A set of synthetic questions are then generated on these atoms (with the chunk as the context). Dense retrieval involves finding the closest set of synthetic questions, and associated chunks, to the user query. It is found that retrieval with the atoms leads to higher recall than retrieval with chunks. Further performance gain is observed with retrieval using the synthetic questions generated over the atoms. Higher recall at the retrieval step enables higher performance of the enterprise LLM using the RAG pipeline.

## introduction since the popularized chatgpt as an instructionfinetuned large language model (llm) deployed at scale to the lay market, there has been a substantial uptake on the interest of businesses to incorporate llms in their products for a variety of downstream tasks @xcite @xcite . for most companies, they are interested in using such models as enterprise llms where the model can handle queries related to proprietary on-premise data. it has been repeatedly demonstrated that these llms have general (public) knowledge implicitly embedded in their parametric memory which can be extracted upon querying @xcite . however, the llms do not have implicit knowledge about a specific enterprise's textual database in a custom domain and hence are prone to hallucinate in such situations @xcite . additionally, the transformer-based @xcite llms typically have a limited context window (due to quadratic order in cost of the attention mechanism), which means information for a specific company to be queried over cannot be directly fed-in as a prompt to the llm. due to limited budget, it is typically not feasible to finetune llms on a specific enterprise's data. in particular, with evolving data from ongoing projects, it is challenging to maintain a constantly updated company-specific llm finetuned on new data without catastrophic forgetting @xcite . to tackle this issue, and with retrieval augmented generation (rag) proposed by @xcite , rag-inspired systems have rapidly become the defacto as a zero-shot solution for enterprise llms. at the essence, there are 2 steps: 1. retrieval and 2. synthesis. documents are split into independent chunks, and a retrieval process is applied to identify the relevant chunks to a given query. the retrieved chunks (which should fit into the context window) with the query are passed as the prompt to the synthesizer llm to get the desired response. currently, the bottleneck for most enterprise llms is the retrieval step, where the correct information is not retrieved for the llm to answer the question @xcite . hence, this work focuses on building upon zero-shot approaches to improve the retrieval step for rag. a potential limitation of the rag set-up is that an embedding model is used to retrieve the relevant chunks efficiently when given a query. each pre-calculated chunk has its corresponding embedding stored in memory, which allows the closes chunks to be retrieved by embedding the incoming query into the same space. however, there is a mismatch in trying to match the space of queries and chunks as each chunk can carry a large amount of information. instead, our work looks to represent each chunk as a set of atomic pieces of information. @xcite introduced atomization of text for improving the assessment of summary consistency. these atoms can be structural (e.g. sentences of a chunk) or unstructured where a set of atoms is generated for any chunk. by embedding the atoms instead of the chunks themselves, the relevant atoms can instead be identified (that correspond to a specific set of chunks) for the posed query in the embedding space. the atomic breakdown of the chunk enables more accurate retrieval. we further identify that even with the atomic embedding representations of the chunk, a given atom and the query do not necessarily best align for retrieval as the former is a statement with a piece of information while the latter is a question about locating a missing piece of information. thus, we propose generating synthetic atomic questions. each atom has a set of questions generated, which in turn are embedded. therefore, the embedded incoming query is used to identify the closest set of atomic questions which in turn point to the relevant set of chunks to be passed to the synthesizer llm in the rag pipeline. as enterprise rag operates over a closed set of documents, the generation of the atoms and corresponding synthetic questions is a one-off cost. similarly, the increased set of embeddings to search over for the closest matches for the query embedding is of less concern given the various very efficient algorithms for embedding search such as faiss @xcite . current information retrieval approaches applied to the rag pipeline look at improving the quality of dense retrieval through generation augmented retrieval (gar), where a query is rewritten for high recall retrieval. however, we focus our attention on representing the chunks more efficiently for retrieval (information retrieval literature explore such approaches -see section 2. the contributions: an exploration of how the retrieval step in the enterprise rag pipeline is improved with structured and unstructured atomic representation of a document chunk and further improvement with the generation of atomic questions. ## related work recently, several works have extended rag @xcite . many approaches finetune the components of the rag pipeline. for example, siri2023 ) jointly train the retriever and the generator for target domains. however, our work focuses on exploring zero-shot solutions as finetuning can be a computationally infeasible procedure for many enterprises. in terms of zero-shot approaches, there have been several extensions proposed. @xcite propose hypothetical document embedding (hyde) where an llm is used to transform the input query into an answer form (hallucinations are acceptable) for improved dense retrieval over the chunks. similarly, @xcite suggest a query expansion approach termed query2doc where an llm is used to expand the query @xcite with a pseudo-generated document, which they demonstrate to be effective for dense retrieval. alternatively, we propose approaches that focus on modifying the knowledge base on which retrieval is performed rather than modifying the user queries as is common in gar @xcite @xcite . @xcite retrieve a superfluous number of chunks during the retrieval step. they then re-rank the retrieved chunks with a re-ranker system to identify the most relevant set. similarly, @xcite propose filco to filter out the retrieved documents as an additional step in the rag pipeline. @xcite explore the zeroshot use of llms as alternatives for traditional re-rankers. @xcite additionally incorporate the re-rank steps with gar in an iterative feedback loop. leveraging the comparative abilities of llms, @xcite propose using pairwise comparisons for the re-ranking of retrieved documents. alternatively, @xcite propose raptor as an iterative technique to pass a summarized context (based on the retrieved documents) to the synthesizer. iter-retgen by @xcite follow a similar iterative summarization strategy with llms. finally, activerag @xcite encourages the synthesizer to consider parametric memory rather than just relying on the set of retrieved documents. @xcite summarize all advanced rag approaches as additional preretrieval or post-retrieval steps. pre-retrieval steps include query routing, query re-writing and query expansion. post-retrieval steps include re-ranking summarization and fusion. the synthetic question retrieval over atomized units from the document set is a form of pre-retrieval that operates on the knowledge store rather than on the user query. hence, our work remains complementary with all forms of post-retrieval rag. see appendix figure 3 . traditionally, retrieval of relevant documents for a given query has been well studied (hambarde and proenca, 2023) with approaches such as bm25 @xcite . in recent years, dense retrieval approaches have dominated as efficient retrieval processes where queries and documents are represented as dense vectors (embeddings) and documents are retrieved based on the similarity between these vectors. semantically meaningful vectors have been possible with the series of regularly updated sentence transformers for generating general purpose embeddings including sentence-bert @xcite , consert @xcite , simcse @xcite , dif-fcse @xcite , sentence-t5 @xcite and e5 @xcite . more recently, there have been a series of more powerful embedding models that adapt instruction-finetuned language models as embedders @xcite @xcite @xcite . therefore, this work restricts exploration to dense retrieval. in recent information retrieval literature, chen et al. ( 2023 ) explore what granularity should be used for retrieval. they introduce the concept of breaking a passage into atomic expressions where each encapsulates a single factoid. @xcite argue a document consists of many diverse details. hence, they propose representing a document using multiple (diverse) embeddings to capture different views of the same content. @xcite investigate for doc2query, a method of expanding the content of a document, how hallucinations can be minimized in the generated queries over a document. our work connects these concepts for specifically generating multiple synthetic questions over atoms in enterprise rag. this work is a bridge between methods explored in information retrieval and the rag community. ## retrieval for rag in enterprise rag systems, the core pipeline can be summarized as follows. here, the focus is on improving the retrieval step of the enterprise rag pipeline. for the scope of the data considered in this work, we assume that the answer to a specific query is present in only one chunk (i.e. there are no unanswerable queries and multiple chunks are not required to deduce the answer to a question). therefore, the retrieval step task can be defined as follows: task let r(q; c) ∈ 0, 1 denote an oracle relevancy function that returns 1 if a chunk, c, contains the answer to the user query q and 0 otherwise. given a set of n chunks, {c} 1:n , and a user query q, retrieve chunk next, we describe the various approaches for the retrieval step of enterprise rag systems. the focus is on zero-shot approaches that can be applied without any training and we assume we have no-cost in accessing the relevancy function. ## results let's take a look first at the all-mpnet-base-v2 embedder for squad. operating at the chunk scale, where the raw text is embedded for dense retrieval, the standard rag achieves a recall of 65.5% with the top 1, which increases to 89.3% when considering the top 5 chunks retrieved. by applying gar with the hyde baseline at the chunk scale, we do not observe gains. as discussed in section 3.3, the text chunk contains several semantic pieces of information while the re-written query remains related to a single semantic piece of information. hence, it is challenging for the hyde approach to improve recall at the chunk scale. by splitting a chunk into structured atoms (sentences), the final rows of let's now consider bipar from in general, for the re-formatted squad dataset, figure 2 presents how the performance varies with the number of synthetically generated questions on the unstructured atoms. for each recall rate (r@1, r@2 and r@5), two profiles are indicated: 1. a random selection of synthetic questions for the atoms of each chunk; 2. an optimally diverse selection of synthetic questions for the atoms of each chunk. the optimally diverse set of questions is selected as follows. a threshold, τ is selected on the pairwise cosine distance. for the full set of atomic questions generated, the pairwise cosine distances of the question embeddings is calculated for each chunk. if any pairwise cosine distance is below τ , one of the questions is purged. the process if repeated until all questions in the remaining set have pairwise cosine distances of their embeddings above τ . by sweeping τ , the total number of synthetic questions across the corpus changes. one can hence expect that a chunk with more information will have a more diverse set of questions. figure 2 shows that a significant number of questions are redundant across the squad and bipar chunks. by removing more than half of the questions (and hence halving the storage cost), performance can be maintained at the maximal value for each of the recall rates. in the extreme setting, with only 20% of the questions retained, there is only a marginal decrease in recall when using the optimal set. thus, despite a larger storage cost with atomic question retrieval compared to standard enterprise rag, the performance boost can be justified with an efficient choice of synthetic questions to retain. see appendix section b.2 for unanswerability analysis of the generated questions. ## conclusions rag systems are a popular framework for enterprises for automated querying over company documents. however, poor recall of relevant chunks with dense retrieval causes errors to propagate to the synthesizer llm. previous works have focused on extensions involving generation augmented retrieval where the query is re-written at inference time to improve recall. conversely, we explore adaptations to the storage of the chunks. the retrieval step for rag can be refined in a zero-shot manner by 1) atomizing the chunks and 2) generating questions on the atoms. significant improvements are observed on the bipar and squad datasets with this approach as partitioning a chunk into atomic pieces of information allows dense retrieval with the query to be more effective. moreover, operating in the question space, the query embedding aligns better with the synthetic questions of the target chunk. we further demonstrate that the storage cost of a large number of synthetic question embeddings can be dramatically reduced by only storing a diverse set of questions for each chunk. question-based retrieval using atomic units will enable the deployment of higher performing enterprise rag systems without relying on any additional training.