LLMs with Retrieval Augmented Generation (RAG)

Introduction

Large Language Models (LLMs) have revolutionized the field of artificial intelligence, enabling applications such as chatbots, virtual assistants, and content generation. However, they face significant challenges, including factual inaccuracies, outdated information, lack of domain-specific knowledge, and the high costs associated with fine-tuning ¹

¹ Finetuning is the process of adapting a pre-trained model to a specific task or dataset. This process can be prohibitively expensive, often requiring substantial computational resources and time.

We are going to first explore the idea of the “context window” of an LLM, and then introduce Retrieval Augmented Generation (RAG) as a solution to these challenges. RAG combines the strengths of LLMs with external knowledge sources, allowing for more accurate, up-to-date, and contextually relevant responses.

The “Context”: How an LLM sees the world

An LLM is trained on a vast corpus of text data, learning to predict the next word in a sequence based on the words that came before it. This training process enables the model to generate coherent and contextually relevant text. What does it mean to be “contextually relevant”?

Looking at a very simple example, consider the interaction with ChatGPT shown in Figure 1. In this brief interaction, we see a user ask about the 2020 world series. ChatGPT responds with the correct answer, showcasing its ability to recall specific information from its training data (and, in this case, it also used web search, but we’ll ignore that for now). In a second interaction, the user asks about the location of the game. After each question, the assistant uses the full “context” to generate the new response.

Figure 1: In this brief interaction with ChatGPT, we see a user ask about the 2020 world series. ChatGPT responds with the correct answer, showcasing its ability to recall specific information from its training data (and, in this case, it also used web search, but we’ll ignore that for now). In a second interaction, the user asks about the location of the game. After each question, the assistant uses the full “context” to generate the new response.

So, a chat session with a large language model is a back-and-forth exchange where each new input builds on the previous context. After each user input, the model runs a new set of predictions, taking into account the entire conversation history up to that point. Of course, the model also has to keep track of who generated which part of the conversation, the order of the messages, and potentially other information such as timestamps or information about the user (e.g., location, preferences, etc.). For the purposes of this discussion, we will ignore these additional complexities.

In Figure 2, we see a simplified representation of how an LLM represents a conversation. The information is displayed in a structured format², showing the user inputs and model responses in sequence. The information in this structured format is the “context” that the model uses to generate its responses.

² The format shown here is called JSON which stands for JavaScript Object Notation. The actual internal representation used by the model is not JSON, but rather a sequence of tokens. However, JSON provides a convenient way to visualize the structure of the conversation.

Figure 2: A representation of the context associated with the interaction in Figure 1. The context is displayed in a structured format, showing the user inputs and model responses in sequence. After each turn of a chat, a new user input and a new assistant response are added to the context.

So, in practice, at each user interaction, the model is not just looking at the most recent user input, but rather the entire conversation history. This is what allows it to generate responses that are relevant to the ongoing discussion.

Figure 3 illustrates how the context builds up over multiple turns in a conversation. Each time the user provides input, the model generates a response based on the full context of the conversation so far.

sequenceDiagram
    participant U as User
    participant M as Model
    Note over U,M: Context: Empty
    U->>M: User Input 1
    Note over U,M: Context: [User Input 1]
    M->>U: Model Response 1
    Note over U,M: Context: [User Input 1, Response 1]
    U->>M: User Input 2
    Note over U,M: Context: [Input 1, Response 1, Input 2]
    M->>U: Model Response 2
    Note over U,M: Context: [Input 1, Response 1, Input 2, Response 2]

Figure 3: A simplified sequence diagram showing how an LLM uses the full conversation history (context) to generate responses. From the model’s perspective, each step in the conversation is an independent event, but the context is built up over several turns.

Returning to the idea of “contextual relevance,” we can see that the model’s ability to generate relevant responses is directly tied to the information contained within this context. And the context is really just a sequence of text, which the model processes to generate its output.

Now that we understand how an LLM uses context to generate responses, we can explore the limitations of this approach. The model’s knowledge is limited to what it has seen during training, and it cannot access new information unless it is explicitly provided in the context.

Take-home message

LLMs do NOT learn from interactions with users EXCEPT through what is contained in the context. The context is the conversation history, which the model uses to generate responses. The model’s knowledge is static and limited to what it has seen during training. It cannot access new information unless it is explicitly provided in the context.

Specific limitations of LLMs

Remember that LLMs are trained on a fixed dataset, which means they have a static knowledge base. They DO NOT learn from new information unless it is explicitly provided in the context. In addition, the training data may not cover all possible topics or may contain inaccuracies. Finally, since the model responses are probabalistic, they may not always be reliable or accurate, particularly for specialized or niche topics. Let’s look at some specific limitations:

Factual inaccuracies and “hallucinations”

Standard LLMs can sometimes generate false or nonsensical information, known as hallucinations, especially when asked about specific, niche, or newly developed topics. RAG mitigates this by “grounding” the AI’s response in verified facts from a controlled knowledge base, which can also provide citations for easy fact-checking.

Outdated information

A conventional LLM’s knowledge is limited to its training data, which becomes outdated over time. RAG allows the model to access current, dynamic data, such as live news feeds or frequently updated product manuals, ensuring responses are relevant and timely.

Lack of domain-specific knowledge

Without RAG, general-purpose LLMs lack the specialized knowledge required for specific contexts, such as an organization’s proprietary data or a technical field like medicine. RAG enables an LLM to access and understand information from a company’s internal documents, customer logs, or scholarly research.

If we think for a moment about these limitations, we might realize that they all stem from a common root cause: the model’s knowledge is static and limited to what it has seen during training. The model cannot access new information unless it is explicitly provided in the context. Why not just update the model’s training data? This brings us to the final limitation:

High cost and resources of fine-tuning

The traditional method of updating an LLM’s knowledge involves fine-tuning or retraining the entire model, which is a computationally expensive and time-consuming process. RAG eliminates this need, allowing developers to inject new data and customize an LLM’s output by simply updating the external knowledge base.

Addressing LLM limitations

Before jumping into Retrieval Augmented Generation (RAG) in detail, let’s look at what some commercial LLMs are doing to address these limitations. For example, ChatGPT has the ability to use web search to access current information. In Figure 4, we see an example interaction where the model retrieves several relevant resources before generating a response about recommended dosing for our patient.

Figure 4: ChatGPT, like most commercial LLMs, has the ability can use web search to access current information. In an example interaction, the model retrieves several relevant resources before answering about recommended dosing for a 72-year-old male undergoing knee replacement.

A few questions come to mind.

• How does the model decide that it needs to use web search?
• How does the model decide which resources to retrieve?
• How does it use the retrieved information to generate a response?

Let’s break this down.

How does the model decide that it needs to use web search?

In a typical interaction with an LLM, the model generates a response based on the context of the conversation. If the model determines that it does not have sufficient information in its training data to answer a question accurately, it may decide to use web search to retrieve additional information. This decision is typically based on the model’s internal confidence level in its ability to generate a relevant and accurate response. If the model’s confidence is low, it may trigger a web search to gather more information.

Chatbots like ChatGPT actually undergo additional training to respond to user queries. This process is often called Instruction Tuning. During this process, the model is trained on a dataset of user queries and corresponding responses, which may include examples of when to use web search. This helps the model learn to recognize when it needs to access external information to provide a more accurate response. In addition, the model may be fine-tuned to recognize specific keywords or phrases that indicate a need for web search, such as “current,” “latest,” or “recent.” Finally, the model may also be trained to recognize specific topics or domains where web search is likely to be necessary, such as medical information or current events.

In other words, the model MUST BE TRAINED to recognize when it needs to access external information to provide a more accurate response.

How does the model decide which resources to retrieve?

When the model decides to use web search, it typically generates a query based on the user’s input and the context of the conversation. This query is then sent to a search engine, which returns a list of relevant resources. The model may use various techniques to rank these resources based on their relevance to the user’s query, such as keyword matching, semantic similarity, or other relevance metrics. The model may also use additional information from the context of the conversation to help determine which resources are most relevant. For example, if the user has previously mentioned a specific topic or domain, the model may prioritize resources that are related to that topic.

How does it use the retrieved information to generate a response?

We are back to the idea of “context.” In a standard LLM interaction, the context is limited to the conversation history. When the model retrieves information from web search, it can incorporate this information into the context of the conversation. This means that when the model generates a response, it can draw on both the conversation history and the retrieved information. The retrieved information is typically added to the context in a structured format, such as a list of URLs or a summary of the key points from each resource. The model can then use this information to generate a more accurate and relevant response.

Retrieval Augmented Generation (RAG)

Retrieval Augmented Generation (RAG) is a technique that enhances the capabilities of large language models by integrating them with external knowledge sources.

A RAG system typically follows a multi-stage process to answer a user’s query:

Indexing

A knowledge base (e.g., internal documents, websites, PDFs) is prepared by breaking down the content into smaller “chunks” as shown in Figure 5. An embedding model converts these text chunks into numerical vectors, which are then stored in a system that will allow for efficient searches.

Figure 5: Indexing a document by breaking it into chunks, converting the chunks into vectors using an embedding model, and storing the vectors in a vector database. When a user submits a query, it is converted into a vector, and a semantic search is performed to find the most relevant document chunks. The retrieved information is added to the context, and the augmented prompt is fed into the LLM to generate a response.

However, the index could also be based on patient ID, medical record number, or some other identifier. The key point is that the index allows for efficient retrieval of relevant information based on a user’s query.

Querying and Retrieval

When a user submits a query, it is converted into a vector using the same embedding model. The system then performs a semantic search to find the most relevant document chunks from the database. This process is illustrated in Figure 5. The most relevant chunks are retrieved and prepared for inclusion in the context.

Augmentation and Generation

The final step involves two sub-steps: * Augmentation : The retrieved, relevant information is added to the context, potentially with some additional processing to ensure it is in a suitable format and with prompt engineering. * Generation : This augmented prompt is fed into the large language model, which generates an accurate and contextualized response based on the new information.

Figure 6: An overview of the Retrieval Augmented Generation (RAG) process. The process involves indexing a knowledge base, querying and retrieving relevant information, and augmenting the context for the LLM to generate a response.

The overall RAG process is illustrated in Figure 6.

Benefits of RAG

RAG offers several advantages over traditional LLMs:

• Enhanced accuracy and relevance
• Access to up-to-date information
• Domain-specific knowledge
• Cost-effective updates
• Improved user trust (via citations)
• Can leverage private or proprietary data
• Can support multi-modal data
• Can use existing databases and query patterns to retrieve information (e.g., medical records, financial data, etc.)

Example: Google’s NotebookLM

Google’s NotebookLM is a RAG-based application that allows users to create a personal AI assistant by uploading their own documents, such as PDFs, text files, and web pages. The system indexes the content of these documents and enables users to ask questions and receive answers based on the information contained within their uploaded files. In addition, NotebookLM has add-on capabilities that are designed to help with learning, including audio and video summaries, flashcards. It can also generate blog posts, informational pamphlets, and other content based on the user’s documents.

I have created a “Notebook” in NotebookLM that contains a few documents related to physician and patient perspectives on the use of AI in healthcare. In Figure 7, we see a screenshot of the NotebookLM interface showing the uploaded documents.

Figure 7: A screenshot of the NotebookLM interface showing the uploaded documents.

The shared notebook can be found here. After clicking the link, you will need to sign in with a Google account. You can then explore the notebook and see how NotebookLM uses RAG to answer questions based on the uploaded documents.

You can also create your own notebook by going to notebooklm.google.com and uploading your own documents. You can learn more about NotebookLM and how to use it by visiting the official website.

Conclusion

Retrieval Augmented Generation (RAG) is a powerful technique that enhances the capabilities of large language models by integrating them with external knowledge sources. By addressing the limitations of traditional LLMs, RAG enables more accurate, up-to-date, and contextually relevant responses. Applications like Google’s NotebookLM demonstrate the potential of RAG to create personalized AI assistants that can leverage private or proprietary data. NotebookLM is just one example of how RAG can be applied, and we can expect to see many more innovative applications in the future.