A multi-dimensional retrieval framework for enterprise RAG systems

1. Summary

Retrieval Augmented Generation (RAG), originally introduced in 2020 [1], has become the most popular application pattern with Large Language Models (LLM). Initial implementations of RAG focused on ingesting a few text documents like technical papers and training manuals to answer questions using semantic search. Enterprises have a diverse set of unstructured data assets and hence were motivated to build RAG systems leveraging these assets. Enterprises typically have their data assets spread across multiple types of data stores and formats. These could be relational databases, document hubs, digital workspaces, ticketing systems, productivity applications, emails and chats. As the scale, variety and change velocity of these assets increase, semantic search runs into limitations. To provide effective retrieval in enterprise use cases, semantic search needs to be complemented with other techniques. Based on our experiments and experience in building enterprise-grade RAG, we propose a multi-dimensional approach to retrieval in enterprise RAG systems. Our evaluations demonstrate that by combining searches across multiple dimensions, retrieval performance improves to meet enterprise grade expectations.

2. Enterprise data characteristics

A lot of initial research in RAG systems focused on ingesting a few documents and being able to wriggle the best retrieval performance out of them. The focus is generally on chunk sizes, embedding models and semantic search patterns. In enterprise data, we are dealing with document hubs, digital workspaces, ticketing systems, relational databases, logs and emails. A single RAG application may need to use multiple such data sources to answer questions. When it comes to building retrieval on enterprise data, a few distinct characteristics stand out, that impacts retrieval performance

· Format varies across unstructured, semi-structured and structured data

· Data volume can run into many Gigabytes to a few Terabytes

· Change is constant. Documents are added and updated every day.

· Enterprise information deals with specific topics. These topics could be about specific products, people, projects, customers and events.

· User queries (prompts) in enterprise RAG typically focus on a topic and profiling information about a topic. A query “When is the delivery date for Project Alpha?” focuses on the topic “Alpha” and the profile “delivery date” about that topic.

· Multiple documents will cover the same topic. The same document can cover multiple topics. They have many cross references to other documents and topics

· Multiple versions of (contradicting) information may exist about the same topic as the topic evolves over time. This brings in the need to keep track of the latest versions over time

· Documents and associated information lose relevancy over time. A log file may lose relevancy in 15 days, a project status report may be valid only for a month, until the next month’s report is available.

· User queries in RAG system may require the latest information only or information from a specific date and time range. A user query “What is the current plan for product X?” needs to filter for the latest information about X. A query “Who worked on project Y in June 2023” will need a date filter for the user requested date range.

· Enterprise data sources have well enforced access controls. A document hub like SharePoint or Google Drive may have 100 users, with each user having access to a unique subset of documents in that hub. When this document hub is indexed and all the 100 users are provided access to the resulting RAG application, each user should only be able to see information retrieved from their own allowed subset of documents.

3. Dimensions for enterprise retrieval

Retrieval systems for enterprise RAG needs to focus on four dimensions to provide secure access to relevant indexed data.

· Semantic dimension focuses on semantic similarity between the user query and the indexed data.

· Object dimension deals with topics (people, projects, products etc.) and profiles (what information is needed about topics). It focuses on matching topics and profiles between the user query and the documents in the index.

· Time dimension focuses on the temporal aspect of retrieval. It filters relevant documents based on their associated timestamps.

· The authorization dimension ensures that enterprise users are provided answers only from those documents they have access to, in their original data sources.

Using these dimensions together will help build an enterprise system that provides secure access to relevant data. These dimensions can be combined with other retrieval patterns in a complimentary manner to provide effective retrieval. In this paper, we will discuss how the Semantics-Objects-Time-Authorization framework (SOTA for short) can be used for building effective enterprise grade retrieval systems.

4. Semantics

Semantic search has been the core driver for the original RAG systems since the advent of LLMs [2]. This technique uses embeddings to search a vector database for documents that are semantically similar to the user query. There is enough literature on several techniques applied in semantic search, so we will not delve into them in this paper. However, this technique also has several shortcomings, especially when applied to the enterprise context. Enterprise use cases have their own vocabulary, containing product names, people names, project names, abbreviations and techniques that are unique to the enterprise. An embedding model built on public datasets does not capture the meaning and relationships of these unique terms [3]. Using such an embedding model for enterprise RAG does not incorporate these special terms. Semantic search using these models will not provide enough attention to the enterprise terms and could lead to lower accuracy levels. Worse, when enterprises reuse general terms for special reasons, this reuse is also not modeled. For example, an enterprise can call an internal IT project “Diablo”. Public embedding models will associate this term with volcanoes, not with the IT project nor its purpose. Building custom embedding models for each use case is also cost prohibitive. Semantic search does not factor in for the temporal aspects of user queries. Semantic search also does not consider any authorization or access control constraints.

5. Objects

Every enterprise has its own set of objects. Objects are of two types. The topic is what the query is about. Topics include people names, project names, product names, abbreviations and other terms that are used within the enterprise. Profile is what information is asked about the subject. This includes descriptions, metrics, timelines, relationships and other types of information. Data assets in the enterprise are filled with these objects. Search queries almost always contain an object, about which the user is asking for information. They may also have an explicit profile, the information that’s requested about the subject. A query “When is the delivery date for Project Alpha?” focuses on the topic “Alpha” and the profile “delivery date” about that topic. Identifying these objects in the user query and providing special consideration for these objects will improve search results considerably.

5.1. Identifying Objects

Extracting objects from a query is difficult and critical. Several techniques exist for doing this extraction, each with their own advantages and shortcomings

5.1.1. Parts of Speech extraction

An NLP library like spacy [4] can be used to extract parts of speech from a given user query. Proper nouns, compound nouns, direct objects, preposition objects and nominal subjects can be potential objects. Using parts-of-speech extract is fast and cheap but is limited to the capabilities provided by the library. It can also extract too many objects, leading to searches returning no results.

5.1.2. LLM based extraction

An LLM can be used [5] to understand the query and extract objects out of it. It will be a sophisticated extraction option but adds significant costs to query execution.

5.1.3. Enterprise dictionary

An enterprise dictionary can be created to contain all the topics and profiles used in the enterprise [6]. The query can be compared with the dictionary to extract known topics and profiles from it. While this option improves performance, it’s an expensive operation to keep the dictionary up to date.

5.2. Using Objects

Objects identified from the query can then be used during the search process. Full-text search is a great option to identify documents that contain the list of objects in the query. Several options exist here for search. An “AND” search can look for documents that match all the objects. An “OR” search can look for documents that contain any of the objects. Alternatively, each document can be scored based on the number of objects matched in the document. This score can then be used for filtering or sorting.

6. Time dimension

The date and time dimension plays a key role in all enterprise data assets. Each event, record or document is typically labeled with a few time attributes like creation time, last modified time, last accessed time etc. In addition, time elements exist within documents also, like due dates. Information about objects is constantly changing in an enterprise. Records and documents are constantly created, updated and archived. When search indexes are built, they will over time accumulate multiple versions of documents. Sometimes different documents may contain the same information, but valid for different points in time. All search queries implicitly ask for the “current” or “latest” information about the object. This is the “recency” aspect of information. Sometimes, search queries explicitly ask for information for a given date range.

6.1. Adding time metadata to indexes

When records or documents are fetched from data sources, typically time metadata is available as created-timestamp or last-modified-timestamp. This should be used as time metadata during indexing. This is straightforward.

6.2. Recency scoring

When searching either with semantic or full-text search, several documents may appear to have information about a given topic and profile. For example, several copies of a project plan may exist with the same tasks, but with different due dates. When performing search, the latest document needs to be used to fetch the right answer. In many searches, this needs to be combined with semantic similarity scores, which makes the process cumbersome. To help, we can use a decay function to provide a recency score to each document. This score can then be used along with a similarity score to extract top k results.

To compute decay, a time unit needs to be identified. The time unit can be hour, day, month or year. So, if the unit is identified as month, all documents updated in the same month are considered “equal” for recency purposes. The right time unit varies by each use case. For a real time log analytics search, time unit might be every-5-minutes. For a near-static guideline document, it can be every 6 months.

The decay score can be in the range of 1 to 0, with 1 being the most recent. The older the document the older the score is. An overall time range also needs to be set, within which decay happens. For example, if the decay unit is a month, overall time range can be 2 years. Documents older than that time range will have a decay score of 0. For purposes of decay, two options are available.

A Linear decay function decays the same amount for each time unit. So, if the time unit is a month and the overall time range is 2 years, a document will decay by 0.042 for each passing month. This function is useful for near-static data sources like policies where documents stay relevant for a long time.

An exponential decay function uses an exponential function [7] to decay quickly and take documents out of contention as they grow older. This function is useful for dynamic data sources like system logs, where old logs quickly lose meaning as far as recency is concerned.

6.3. Temporal filters

In addition to recency scores, temporal filters restrict the search to a smaller meaningful subset of documents. This will improve the effectiveness of semantic and object dimensions. There are a few options that are applicable here.

6.3.1. Default date range

A default date range is imposed for the query based on a general range in which relevant data is expected to be available. The range is worked backwards from the current date. The default range is based on the use case. An intranet source may have a default range of 2 years, while a log analytics system may have a default range of 15 days (about 2 weeks). While the index may have data for a longer period, the default date range restricts it to a meaningful window.

6.3.2. Latest date range

Sometimes, queries explicitly ask for the latest or current information. A query like “what is the latest status for project X? “ . In this case, we replace the default date range with the “latest date range”. This range is specific to a use case. The input query is analyzed (based on keywords & semantics) to see if it looks for the latest information and if so, the latest date range is applied. The latest date range could be last 6 months for an intranet data source or last 1 hour for a real time log analytics use case.

6.3.3. User requested date range

If the query asks for a specific date range like “What milestones were completed for project X in June 2023?”, then that date range (for that example 6/1/23 to 6/30/23) will applied instead. The specific date range needs to be extracted from the query either using NLP tools or with the help of an LLM.

7. Authorization

Most enterprise data sources enforce some kind of access control. This is usually tied to the user’s enterprise login credentials. Similarly, enterprise RAG systems also enforce user access based on their enterprise login. A human resources virtual assistant may allow access to all employees, where each employee has access to all the data. A SharePoint used by the marketing team will only allow access for those employees who work for the marketing team [8]. Even there, each employee may only have access to a subset of documents, based on their role. When a RAG system is built for this marketing SharePoint, the same access control at the source needs to be enforced, so the users can only get answers from those documents they have access to, at the source. This is a key requirement for any assistants based on enterprise data.

To enforce authorization, there are two possible modes.

7.1. Pre-ingest access control lists

When documents are ingested from the source system into the semantic index, their access control lists are also ingested and kept as metadata. When performing a search, the credentials of the user who issued the query are used to check this metadata to see if they have access to that document. Pre-ingestion makes the retrieval task simple and quick. However, as access control lists have several roles and layers, ingesting and storing them can be a tedious task. When the access changes for a document, catching up to the change can be a problem. Each data source system may implement access control differently, so trying to create a common system across all data sources will also be a challenge.

7.2. On-demand authorization checks

In this method, no access control lists are ingested. Instead, for documents retrieved from search, an authorization check query is issued to the source system (e.g. SharePoint) to see if the user has access to the document. If not, that document is removed from the list of results. If few or no documents are found, the search is repeated to retrieve additional documents. This method ensures that the authorization is always current. But it leads to a complex retrieval process and increases latency for answering the query.

8. Implementing multi-dimensional retrieval

When multiple dimensions are used for retrieval, they need to be combined in a complimentary fashion to create the retrieval query. To add dimensions to a query, there are two popular methods.

Filters

Filters are typically implemented in the “WHERE” clause of SQL queries. There are similar equivalent constructs in other retrieval systems. Filters can either be actual values (e.g. objects, dates) or on scores (similarity, recency). Filters restrict the records received at the source. Combining too many filters can prove to be too restrictive. This can lead to zero documents or fewer than the number of results expected.

Sort

In the case of sorting, the retrieved documents are sorted by specific attributes or computed scores. Popular sorting elements include dates, similarity scores and recency scores

The following table 1 provides the list of techniques in SOTA and whether filtering or sorting are applicable to them

Table 1:

Which dimensions to use for filters and which ones for sorting? This decision is specific to the use case. Not all dimensions are applicable for all use cases. First, the list of dimensions needs to be decided for the use case. Then, experimentation needs to be done by trying various combinations for filtering and sorting. The combination that provides the best outcomes would be the one that is chosen.

9. Progressive relaxation

Using all the four SOTA dimensions in a query can lead to an ultra-restrictive search, leading to few or no documents in the results. They can also restrict semantic search from exercising its full potential, which is the key driver for RAG. To overcome this, a progressive relaxation technique can be applied. In this case, we can remove filters one by one until enough documents are retrieved for generation. Authorization is one dimension that cannot be compromised for enterprises, so that dimension cannot be relaxed. Object and Time dimensions provide the levers for relaxation. A progressive relaxation algorithm can work like this. This is just an example, and multiple combinations can be tried based on the use case and priorities for the dimensions

Step 1: Run query with Semantics, Objects, Time & Authorization dimensions to get k documents

Step 2: If Step 1 returned no documents or fewer documents than k, remove the time dimension and rerun

Step 3: If Step 2 returned no or fewer documents than k, additionally remove the object dimension and rerun.

10. An example evaluation

To evaluate the effectiveness of SOT a sample dataset was used. This dataset contains a list of published articles about a company’s products, services and interviews. The articles are also timestamped by the date of publication. When users execute queries, the retrieval system should pull up the latest articles for the topics mentioned in the query. So, in addition to the semantic dimension, the dataset also has the object dimension (products, services, people) and time dimension (recency). Authorization was ignored for evaluation purposes. The retrieval system uses objects and date ranges for filters. For sorting, it uses semantic similarity and recency. Progressive relaxation is also used to remove recency and objects iteratively.

The labeled evaluation dataset has 100 queries and has only 1 correct result document for each query. During evaluation, the position at which the correct document is in the list of retrieved documents is measured. Using this, we measure recall at each of the positions from 1 to 10. So, if the recall is 50% at k=3, it means that the right document is returned within the top 3 results for 50% of the queries.

For the purposes of evaluation, three types of tests are run. This is to measure the additional performance gain by adding an additional dimension.

1. Semantics, Objects & Time (with progressive relaxation)

2. Semantics & Objects only (with progressive relaxation)

3. Semantics only.

The results are plotted in image 1

Image 1:

In the results, we see that the plain vanilla semantic search performs much lower than the other two approaches. At k=5, semantics only search has recall at 65%. It improves to 85% when objects are added. It improves to 89% when temporal dimension is added. The additional gains for objects and temporal dimensions will vary based on the use case.

11. References

[1] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, Douwe Kiela. Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. arXiv preprint arXiv: 2005.11401, 2020.

[2] Bast, Hannah; Buchhold, Björn; Haussmann, Elmar (2016). “Semantic search on text and knowledge bases”. Foundations and Trends in Information Retrieval. 10 (2–3): 119–271. doi:10.1561/1500000032.

[3] Luke Merrick, Danmei Xu, Gaurav Nuti, Daniel Campos. Arctic-Embed: Scalable, Efficient, and Accurate Text Embedding Models. arXiv:2405.05374v1, 2024

[4] spaCy’s NER model: https://spacy.io/universe/project/video-spacys-ner-model

[5] Cobus Greyling: Using a large language model for entity extraction : https://cobusgreyling.medium.com/using-a-large-language-model-for-entity-extraction-6fffb988eb15

[6] Jin Wang, Chunbin Lin, Mingda Li, Carlo Zaniolo : Boosting approximate dictionary-based entity extraction with synonyms, https://www.sciencedirect.com/science/article/abs/pii/S0020025520303303, 2020

[7] Exponential decay https://en.wikipedia.org/wiki/Exponential_decay

[8] Understand groups and permissions on a SharePoint site https://support.microsoft.com/en-us/office/understand-groups-and-permissions-on-a-sharepoint-site-258e5f33-1b5a-4766-a503-d86655cf950d