What Is Retrieval-Augmented Generation (RAG) in AI?

Retrieval-augmented generation, usually shortened to RAG, is a way of making an AI language model check a body of trusted documents before it answers. The model's reply is then based partly on what was found, not only on what it absorbed during training. It is one of the main techniques used to make chatbots answer questions about current or private information. It helps, but it does not make a model infallible.

The problem RAG tries to solve

A large language model learns from a fixed set of training data. Anything that happened after that data was collected is unknown to it, and so is anything private, such as a company's internal policies. Retraining a model to add such knowledge is expensive. Amazon Web Services describes RAG as a cost-effective way to improve a model's output without retraining it.

The idea also has an academic origin. In May 2020 a team led by Patrick Lewis, with colleagues from Facebook AI Research, University College London and New York University, posted a paper called Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. It was later accepted at the NeurIPS 2020 conference. The authors noted that large language models keep facts inside their parameters, but find it hard to access and precisely use that knowledge, and that they trailed task-specific systems on knowledge-heavy tasks.

How RAG works, step by step

AWS lays out four parts to a RAG system:

• Prepare the data. Documents from databases, APIs or file stores are converted into embeddings, which are lists of numbers that capture meaning, and saved in a vector database.
• Retrieve. The user's question is also converted into an embedding and compared with those stored. The closest matches are fetched. In AWS's example, a question about annual leave brings back the leave policy and the employee's own leave record.
• Augment the prompt. The retrieved text is added to the question as extra context, and the model is asked to answer using it.
• Keep the data fresh. Documents and their embeddings are updated, either continuously or in batches, so answers do not go stale.

The 2020 paper used a pretrained text-to-text model paired with a dense index of Wikipedia, searched by a neural retriever. The authors called the first part parametric memory, because it lives in the model's parameters, and the index non-parametric memory, because it sits outside the model and can be changed.

What it is good for

According to AWS, the main benefits are these:

• Lower cost, because the base model is not retrained.
• More current answers, because the system can connect to frequently updated sources.
• Easier checking, because answers can cite the documents they used and readers can open them.
• More control for developers, who can change the sources and limit retrieval by who is allowed to see what.

In the research paper, RAG set state-of-the-art results on three open-domain question answering tasks at the time. The authors also reported that its generated text was more specific, more varied and better grounded in facts than that of a baseline model without retrieval.

Where it falls short

RAG moves part of the problem rather than removing it. If the retrieval step finds the wrong passages, the model has poor material to work with. AWS notes that finding the right context gets harder at scale and that keyword search alone gives limited results for knowledge-heavy tasks. Preparing the data also takes work: documents have to be split into sensible chunks and converted, and the stored copies must be kept up to date.

AWS describes RAG as one approach to some of the challenges of language models, not a full cure. The model can still misread what it is given or add details of its own. The AWS page gives no accuracy or speed figures for RAG as a whole, so any claim of a specific improvement should come from a test of the actual system.

Where you may meet it

Many chat tools that answer from company manuals, help centres or uploaded files use some form of retrieval before generating a reply. Cloud providers sell managed versions; AWS names Amazon Bedrock knowledge bases, Amazon Kendra and Amazon SageMaker JumpStart as examples. Other vendors offer their own. This article is a general explanation and does not compare products.

Sources

• What is RAG (Retrieval-Augmented Generation)?, Amazon Web Services
• Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks, arXiv (Lewis et al., 2020)