The big ideas behind retrieval augmented generation

LLMs are trained on generic public data that’s often dated by the time the chatbot can use it. They work as well as they do because there is a tremendous wealth of such data available today, but what about the data they don’t have access to? My daughter’s insistence that we do her homework according to the precise method they’ve taught in school may seem like a trivial example (though, not for her!). Yet, parallel scenarios are playing out for users who want to turn this new technology into a product they can sell but who need the LLM to be able to talk about their private content. 

Additionally, LLMs are prone to hallucinating if they don’t have the answer to a question in their training set. They will often confidently tell users completely untrue information to provide something instead of a disappointing "I don’t know." There are things these models shouldn’t know, though. We don’t want them to have access to our proprietary information, and we certainly don’t want them to make up answers to questions that can only be answered using that proprietary information.

So, are LLM-powered chatbots destined to stay amusing distractions forever and of no actual use in production and at scale? Of course not! Let’s discuss how we can get the best of both worlds: lovely natural language answers grounded in information from private data sources.

The first step to achieving our goal of improving generative AI with information it doesn’t have access to is to identify what information is needed. How good does that search need to be, though? Assume that you have a data set that includes all of the answers to the questions you want to ask your chatbot. Let’s consider a few possibilities and some case studies for each.

This option would be ideal in a perfect world, but the realities of training an LLM make this approach impractical for most companies. LLM training requires truly massive amounts of high-quality data on the order of hundreds of billions of tokens¹ as well as well-trained data scientists working on expensive computing resources. 

One company that was able to achieve this is Bloomberg with BloombergGPT, an LLM that is designed to work well with finance-specific content.² To do this, Bloomberg used its 40 years’ worth of financial data, news, and documents and then added a large volume of data from financial filings and generic information from the internet. Bloomberg’s data scientists used 700 billion tokens and 1.3 million hours of graphics processing unit (GPU) time. Most companies simply don’t have these kinds of resources to play with. So, even if your company has all the data to answer questions correctly, it’s likely far from enough information to train an LLM from scratch.

This second option requires much less computing power and time than training a model from scratch, but this approach still has significant costs. Google took this approach with Med-PaLM, an LLM for the medical domain.³ What it accomplished is impressive, but we can’t all be Google. The work required a large team of data scientists as well as an LLM Google was allowed to fine-tune in the first place. Not all LLM vendors allow fine-tuning, such as OpenAI with GPT-4.

Fortunately, another approach works within the parameters established by existing LLM providers, taking advantage of the enormous effort already put into LLM development. This technique, RAG, uses prompt engineering to allow existing LLMs to work with additional content without modifying the LLM simply by including the necessary information in the prompt. A good example of this approach in action is the Elastic Support Assistant, a chatbot that can answer questions about Elastic products using Elastic’s support knowledge library. By employing RAG with this knowledge base, the support assistant will always be able to use the latest information about Elastic products, even if the underlying LLM hasn’t been trained on newly added features. 

To make this approach practical and scalable, we must carefully identify what this extra content should be. LLM vendors typically charge by the token, so every word sent to them counts. The tables below show some of those costs for several popular LLMs.

So, as tempting as it might be to just send the LLM everything in your data set, this simply isn’t a practical approach to prompt engineering. Instead, we want to come up with the tightest possible prompt to help the LLM answer correctly. We’ll discuss the pieces that go into constructing a prompt later, but first, we need to figure out how to find the additional content we want to include. How do we identify what the LLM will need to give the correct answer?

Text search has long been dependent on keyword-focused algorithms. Keyword search is a syntactic approach to information retrieval. In linguistics, syntax refers to how words and phrases are arranged together. It doesn’t matter what the words mean. Consider the following sentences:

1. Jack and Jill went up a hill.

2. A hill is what Jack and Jill went up.

3. Jack and Jill went up a dune.

These sentences all mean the same thing, but only the first two are matches using keyword search on the word hill. Keyword search is a well-understood problem and works exceptionally well to a point. A search algorithm using keywords will never return sentence (3) given the search term hill. Getting sentences (1) and (2) may be all that is required for the use case, but if we want sentence (3) as well, we need to move beyond syntax to semantics.

Semantics refers to the meanings of words and phrases in linguistics. Computers find this much more challenging to “get” as it requires a depth of understanding that is difficult to program. Average English speakers probably know that hills and dunes are similar concepts, but how do we teach this to computers? Well, one thing that computers are really good at is math. If we can turn the problem of semantic search into a math problem, computers will be great at it! So, that’s what we’ll do.

Let’s begin with a simple example. We want to describe the meaning of a concept in a numeric way. Imagine we are describing the concept of a cat. For our example, we’ll use images to represent concepts. We could have a scale from –1 to 1 that says how much a given image looks like a cat. In Figure 1, we represent this on a number line, which you can think of as the x-axis of a coordinate system. The picture on the far right is assigned the value [1] since it looks very much like a cat. The airplane on the left definitely does not look like a cat, so it gets the value [–1]. The image at point [0] has some cat-like features but not enough to score higher on the looks-like-a-cat scale.

Since concepts are located in the model based on specific features, concepts that are near each other in the model are likely similar in some way. These nearest neighbors are what we return when we’re doing semantic search. To find similar concepts, we first locate the concept we’re looking for in the model. Then we see what other concepts are nearby. Very close concepts are better matches for the search term, while more distant neighbors are not as good. Figure 4 shows a simple example of this using our two-dimensional vector space or model.

In the old days, when you were trying to answer a question using the enormous amounts of data on the internet, you would carefully consider your keyword selection and hope that one of the returned links contained the answer you needed. As search engines have improved, they’ve started to include a preview of where they think the answer is on the page, so you sometimes don’t even need to click on the link. It’s still a bit of a scavenger hunt to get a complete answer, so it is very compelling to have generative AI chatbots bring all the pieces together for you. 

To accomplish this, we first need to find the pieces. Essentially, we are doing what we needed to do in the old days and then having the chatbot figure out the complete answer from those “links.” In the previous section, we introduced semantic search as an alternative to syntactic search to improve the relevance of what we send to the chatbot. It should be clear now why searching for meaning can provide better results than finding exact matches. 

What is the meaning of the content on the other end of a hyperlink, though? You’ve probably experienced following a link and then getting lost on the page trying to figure out why that link was returned in the first place. There are likely hundreds, if not thousands, of tokens on the page. We already discussed why it isn’t practical to send all of the tokens to the chatbot. It can be costly to send so much data, and there are also technical limitations on how much data you can send at a time. So, it doesn’t make sense to find the “links” and send everything on the page to the LLM. We need to break the content into meaningful pieces before we perform semantic search.

This process of breaking the content up is called chunking, and it is necessary for reasons beyond the limitations of what you can send to the LLM. Consider the example of the mischievous cat about to pounce again. A vector embedding is a numerical representation of a concept, but there are at least four distinct concepts in this phrase.  We don’t really care about the distinct concepts, though. We care about them taken all together, so what we really want is a single embedding for the entire phrase. 

This is something today’s embedding models can do, though the details of how they do this are beyond the scope of this post. In simple terms, the embeddings for mischievous, pounce, about, and to influence the numerical representation for cat, allowing us to perform semantic search over complex concepts like the example. Yet, consider how it would work to assign a single embedding to an entire paragraph about our mischievous cat. What about a short story or even an entire novel?

You naturally lose some of the detail as you include more concepts in the vector embedding. That is, semantic precision goes down as you include more content. For example, a novel can be about many things, not just a single concept. On the other hand, you’re pretty much guaranteed to find the “answer” to your question if you send the entire novel to the LLM. We know we can’t realistically do that, but there is another reason why we can’t vectorize an entire novel. Just as LLMs limit how much data they can receive at a time, most embedding models limit how many tokens can be included in the embedding. After that token limit is reached, the rest of the content is simply dropped. So, even if we try to get a single embedding for a novel, we’ll probably only capture the meaning of the first paragraph or so and then ignore the rest.

So, the task of chunking the data is trickier than it might first appear. As shown in Figure 5, the size of the chunks critically influences our ability to uncover accurate answers. Fine grained or small chunks have high semantic precision since the vector representation for the chunk is less diluted by multiple concepts. However, small chunks are less likely to contain the complete “answer” we’re looking for. Imagine if every link Google returns only points to a few words. You’ll almost certainly need to click through multiple links to get all of the information you need. On the other hand, coarse grained or big chunks likely have complete answers, but it’s harder to identify them as good matches since the numerical representation is influenced by so many concepts.

System prompt
First, there is always a system prompt that defines how we want the chatbot to behave in general. With ChatGPT, this system prompt isn’t shown, but when we create our own RAG application, we need to define it in order to get the chatbot to do what we want. Since we will be retrieving data to augment the chatbot’s knowledge, we need to make sure that the generation uses that data. Moreover, if we want to restrict the application to using the retrieved data to prevent hallucinations, this should also be specified in the system prompt. For example, the following system prompt would ensure the behavior we’re looking for:

You are a helpful AI assistant who answers questions using the following supplied context. If you can’t answer the question using this context, say, “I don’t know.”

Supplied context
The next part of the prompt is the supplied context. This is where we include data we retrieved using search. For example, if we run semantic search and find the three nearest neighboring chunks for the search term, we can provide those three chunks in the supplied context. Alternatively, we might only send the closest nearest neighbor along with the preceding and succeeding chunks. It’s up to us what we include in this section, and it can take some trial and error to figure out what works best in our application.

User input
The final piece of the prompt is the user input. When we use a chatbot like ChatGPT, we are restricted to use only this part of the prompt. For our RAG application, however, this section serves a dual purpose. As usual, this is where we ask the chatbot to do something specific, such as explain how to answer a particular algebra problem, but this is also where we supply the search terms we’ll use to provide context to the prompt.  With all of these pieces in place, we’re ready to see our RAG application in action.

Figure 9 shows how this would work for our original problem. How can we get a chatbot to answer questions about algebra using only information provided in class? We first identify what information is permissible and ingest it into a searchable data store, such as Elasticsearch. For example, my daughter’s teacher provided class notes in PDF form as well as a link to an approved site for introductory topics. As we ingest the data, we need to prepare it for semantic search. Fortunately, Elasticsearch makes this extremely easy to do by providing a field type that performs chunking for us and figures out the embeddings for each chunk. We have this easy button for semantic search, and there is also an AI playground where we can see our RAG application work as we develop it. Using the Playground, we can see if our application meets my daughter’s standards (see Figure 10).