Specialize LLM with Retrieval Augmented Generation (RAG)

Overview

Large Language Models (LLMs) have transformed our interactions with AI, showcasing impressive abilities in understanding and generating human-like text. These models are trained on vast amounts of generalized data, excelling at general knowledge tasks and engaging in conversations across diverse topics. However, this generalist nature comes with significant limitations when dealing with specific domain questions requiring access to current, relevant domain information. LLM training data is static and has a knowledge cutoff. Even if the original data sources for LLMs are suitable, it's challenging to maintain relevancy.

LLM Challenges with Specialization

Consider you want to build a customer service AI Agent for a health insurance companies to answer customer's queries such as latest insurance policies, Medical illnesses covered in policies, or Cashless hospitals covered. While an LLM might provide a plausible response based on its training data, it cannot guarantee accuracy of information, potentially leading to errors. One could argue for developing a specialized version of an LLM model for the health insurance company by training it on domain-specific information. However, this approach is not frugal & will come with an additional costs, it would also require extensive computational resources and expertise. Additionally, it poses challenges in keeping the model relevant with frequent changes in information

RAG is one of the approaches to solve these challenges. It redirects AI assistant app to retrieve relevant information from Authoritative, pre-determined knowledge sources and use this as context for LLM to generate response.

Understanding Retrieval Augmented Generation (RAG)

RAG emerges as a solution to enable the development of specialized AI agents without the need for expensive model training or fine-tuning. It extends the powerful capabilities of LLMs to specific domains by combining their reasoning capabilities with a dynamic knowledge retrieval system. Given below let's try to understand the key features of RAG -

Key Features of RAG:

• Allows LLM models to access and reference specific, up-to-date information in real-time
• Optimizes LLM output by referencing an authoritative knowledge base outside its training data
• Enables AI agents to respond with precision using an organization's latest documentation, policies, data sources or domain-specific knowledge

Think of RAG as giving your LLMs a specialized reference library to consult before responding to queries. This specialized library, often referred to as a "Knowledge base," is a collection of various data sources and documentation which the LLM model references before making decisions.

How RAG Works ?

Now that we've understood what is RAG & its ability to create a specialized AI agent. Let's try to understand the working of RAG on a high level. Given below is high level flow of RAG architecture when users query to the AI Model followed by explanation of each flow -

1. Create Knowledge Base: New data outside the LLM's original training dataset is referred to as the knowledge base or external data. Sources can include document repositories, APIs, databases, etc. This data is converted into numerical representations called as vector embedding and stored in a vector database using embedding language models such as Amazon Titan text Embedding.
2. User Query : Whenever user queries, AI Agent first sends query to this embedding model which transform user's question into a mathematical representation (query embedding).
3. Retrieval of Relevant Information - System looks for document chunks with vector embeddings most similar to the query embedding. The embedding model compares query embedding (user's query numerical representation) to vectors embedding of an available knowledge base stored in Vector DB. In short it compares the numerical representation of user's query with numerical representation of knowledge base data. When it finds a match or multiple matches, it retrieves the related data, converts it to human-readable words and passes it back to the LLM.
4. Prompt Augmentation: The LLM model augments the user input by adding relevant retrieved information in context to the user's query. This step uses prompt engineering techniques to effectively instruct the LLM to use retrieved information for response. Consider this as giving exact information to Human from textbook chapter on which the question is asked to help them answer accurately. Example of Augmented Prompt is given below -

// Example of Prompt Augmentation for AI assistant for Health Insurance company
// retrieved_contexts - Knowledge base context retrieved 
def createAugmentedPrompt (query, retrieved_contexts): 
  
prompt = f"""You're a AI assistant for health insurance company. Your job is to      answer the customer's query related to health insurance policies. Answer the question based on the provided context. If the answer cannot be derived from the context, say so. 
 
 Context: {retrieved_contexts}

 Question: {query}
    
 Answer: Let me help you with that based on the available information."""        

Benefits of RAG

Now that we have understood RAG, Let's try to understand few of the benefits which RAG offers on top of LLM

• Specialization: Transforms general-purpose language models into domain experts without expensive training.
• Accuracy: Improves response accuracy by grounding LLM outputs with verified information.
• Relevancy: Maintains up-to-date information by dynamically incorporating the latest data. Update to the knowledge base can be made easily without having to train LLM model
• Cost-Effectiveness: Provides a cost-effective approach to building specialized AI systems as we don't have to train the Model from scratch.
• Enhance Customer Trust: RAG allows LLM to present information with source attribution and increases confidence in AI solutions. Users can look up-to citations or references of source documentations.
• Adaptability & Scalability: Easily adapts to new domains or expands knowledge base without architectural changes

Looking Ahead

In my next blog post, I'll provide a hands-on implementation guide with a real-life example of using RAG with an LLM model. We'll explore technical frameworks that facilitate RAG implementation. Specifically, we'll dive deep into AWS Bedrock, a fully managed service that offers a choice of foundational model and efficient mechanism to create & integrate knowledge base with chosen foundational model.

This upcoming guide will offer practical hands-on lab for developing generative AI solutions using AWS Bedrock.

Stay tuned !!