Retrieval-augmented generation

Retrieval-augmented generation (RAG)¹ is a technique where a search step finds the parts of a document collection that are related to a question, and a language model answers the question using these parts. For example, a support assistant can find the help articles about passwords and answer a question about passwords from them.

The advantage of retrieval-augmented generation over fine-tuning (see Chapter 11) is that a changed document needs no new round of training. In addition, the answer can name the documents that it is based on.

19.1Finding the chunks

Consider a problem where we are given a collection of documents and a question, and our task is to give the language model the parts of the collection that the question needs. A straightforward way is to put the whole collection into the prompt. For example, if the collection has 10,000 articles of 500 words each, every question sends 5,000,000 words to the language model. However, the prompt has a maximum length, and the time and price of each answer grow with the length of the prompt.

The idea is to split the documents into short chunks and to put only the $k$ chunks that are most similar to the question into the prompt. To compare a chunk with the question, we calculate their embeddings with an embedding model and take the similarity of the two embeddings (see Chapter 6).

Real embeddings have hundreds or thousands of dimensions. Suppose that they have only two dimensions, and that the question “How do I change my password?” has the embedding (0.96, 0.28). As an example, consider the following chunks:

For example, the similarity of chunk 2 is $0.8 \cdot 0.96 + 0.6 \cdot 0.28 = 0.936$. The following picture shows the embeddings:

The embeddings of chunks 1 and 2 point in almost the same direction as the question, and the embedding of chunk 4 is perpendicular to it. Thus, when $k = 2$, the prompt contains chunks 1 and 2.

The embeddings of the chunks can be calculated before the first question and stored. After this, each question needs one new embedding and $n$ similarities, where $n$ is the number of chunks. The time complexity of the search is $O(nd)$, where $d$ is the number of dimensions, because each similarity is a sum of $d$ products.

In practice, the embeddings are stored in an index for approximate nearest neighbor search². The index does not compare the question with every chunk, but it can miss some of the $k$ most similar chunks.

19.2Building the prompt

The prompt contains an instruction, the chosen chunks and the question. The prompt for the above example is as follows:

Answer the question using only the following text.
If the text does not contain the answer, say so.

To reset a password, open Settings and choose Security.
A password must have at least 12 characters.

Question: How do I change my password?

From this prompt, the language model can answer that the customer opens Settings, chooses Security and enters a new password of at least 12 characters. Note that the answer needs both chunks. When $k = 1$, the prompt only contains chunk 1, and the answer does not mention the length rule.

A larger $k$ gives the language model more chances to see the chunks that the question needs. On the other hand, the prompt becomes longer, and it also contains chunks that do not help, like chunk 3 when $k = 3$.

The size of the chunks is a similar trade-off. The advantage of short chunks is that each chunk is about one topic, and its embedding is near the questions about that topic. On the other hand, when a fact is split between two short chunks, we can find only one part of it.

19.3Missing chunks

The language model sees only the chosen chunks. For example, the question “How many articles mention refunds?” needs every article, but the prompt contains only $k$ chunks. In this case, the answer counts only the chunks in the prompt. It is too small when more than $k$ chunks mention refunds.

For this reason, the instruction tells the language model to say when the text does not contain the answer. Without it, the language model often answers from its training data, which can be out of date.

In practice, the search often combines embeddings with a word search, because an embedding can miss an exact name or an error code.