Traditional language models, while powerful, are limited by the static nature of their training data—they can't access new information beyond what they learned during training. This is where Retrieval-Augmented Generation (RAG) comes in. Combining external databases (aka knowledge bases) with a generative model enables RAG to use relevant, real-time information to give better responses.
Basics of Knowledge Bases
The data model of a vector database is an index of documents, where documents are composed of an unique id, the embedded data as a vector with a certain dimensionality and optional metadata. The embedded data represents a chunk of information. There are multiple ways to chunk the information for ingestion in the vector database. For instance, PDF documents can be chunked by paragraphs, by page or semantically by using LLMs to determine when the semantics change between sentences. Unlike traditional databases that store data in tabular form and use exact matching for queries, vector databases store data as vectors and use similarity-based search techniques like cosine similarity.
RAG
A classic RAG flow is shown in Figure 1. The first step is the Retrieval of relevant information from the vector database. Therefore, the query
Which are effective pharmaceutical drugs in the treatment of leukemia?
gets embedded by an LLM into the same dimensionality of the embedded documents of the vector database (use one model for both e.g text-embedding-ada-002). A similarity score is calculated between the embedded query and each document. Then, an ordered list of the semantically most relevant documents gets retrieved. The next step is the Augmentation, whereby the documents get added as context to the query, to provide the LLM with all the information it needs to provide an useful answer. The response is then provided in the last step, the Generation step:
Effective drugs for leukemia vary by type: for example, Imatinib for chronic myeloid leukemia, and Brutinib for chronic lymphocytic leukemia.
RAG with Reranking
The resulting response can be significantly improved by using a simple method called Reranking. A RAG flow with Reranking is shown in Figure 3. The query of the user is rephrased before retrieving the documents from the knowledge base into 4 additional queries that are very similar to the original query:
Which are effective pharmaceutical drugs in the treatment of leukemia?
What are the best medications for treating leukemia?
Which drugs are commonly prescribed for leukemia management?
What pharmaceutical treatments are available for leukemia patients?
What are the most effective therapies for leukemia treatment?
Then each query retrieves a list of the semantically most relevant documents, resulting in 5 lists with 5 documents.
Now, to determine which 5 of these in total 25 documents will be the best to answer the first query, a reranking algorithm called Reciprocal Rank Fusion (RRF) is used. Each document is assigned a RRF-score (see Formula below) by taking the reciprocal of the rank in the list (the first element has the rank 1 hence → 1/1, the second 2 → 1/2, the third → 1/3 etc.). A constant k (usually 60) is added to the denominator to mitigate the influence of high ranks.
\(RRF-Score= 1/(rank + k)\)RRF now iterates through all lists and adds the Rank-Scores of a document and updating a list of documents with the highest scores. A schematic explanation is shown in Figure 2, but with only 3 Lists of 3 documents each returning 3 reranked documents.
The resulting answer has much more and accurate information.
Leukemia treatments include Imatinib for chronic myeloid leukemia. Ibrutinib for chronic lymphocytic leukemia, Cytarabine for acute myeloid leukemia, Methotrexate for acute lymphoblastic leukemia, and CAR-T cell therapy for difficult-to-treat cases.
Important things to consider
Reranking is a simple and effective technique, however some things have to be considered for it to work properly.
Embedding quality
The effectiveness of a RAG is mostly determined by the embedding quality of the LLM. A constantly updated list of the best embedding models is maintained on the MTEB leaderboard. If you want to set up a RAG in specific domain like medicine, use a model that was trained on huge amounts of medical data (e.g BioBERT), to make sure that the model embeds medical terms properly.
Enough documents
The larger your knowledge base or the more documents available, the more effectively reranking can improve search relevance.
Rephrasing
Rephrased queries should be similar enough to capture the intent but varied enough to avoid returning identical results.
Thanks for reading!
I hope you learned something! I’ll be posting regularly about interesting things in AI. Stay tuned for more updates and consider subscribing!