Optimizing Retrieval Augmented Generation with Learned Chunking

In collaboration with Elan Markowitz (PhDc)

Introduction

Retrieval Augmented Generation (RAG) techniques recall relevant portions of text to give an LLM context before answering a query. However, this process relies on the documents being preprocessed into smaller, more manageable chunks of text.

Chunking is thus a crucial, but understudied step in the RAG process. The RAG process starts by chunking a corpus of documents into smaller portions. These chunks will then be embedded so that they can be retrieved according to their similarity to a user’s query.

These chunks are preprocessed, embedded with an embedding model, and stored in a database.

These document chunks can then be used during the RAG process. At inference time, the user query is similarly embedded and then the most similar document chunks are retrieved. These chunks are then passed into the context of the LLM for generating its answer.

The LLM that generates answers is optimized for generating text, and the embedding model is optimized for retrieval. However, no method has introduced an optimizable framework for learning the chunking algorithms.

This blog post explores one such approach to learning chunking algorithms.

Motivation

Current text-based splitters generally use rules or heuristic-based approaches to chunking

1. CharacterTextSplitter: Simple Fixed-size character splitter. CharacterTextSplitter is from the LangChain library to split a given text into smaller chunks based on a specific separator and defined chunk size.
2. Recursive Character Splitter: Recursive chunking based on a list of separators. The RecursiveCharacterTextSplitter from LangChain breaks down text by trying different characters (like newlines, periods, commas, and spaces) in a specific order until the chunks are the right size. This method is customizable, allowing users to define the list of characters and the order in which they should be tried for splitting.
3. Semantic Splitting: Here the chunking units are determined based on the underlying meaning(semantics), rather than simply splitting text by sentences or paragraphs. This method first breaks the paragraph into individual sentences. Then, it joins sentences together if they are semantically similar.

Our approach is Learned Chunking. We train a model that predicts where the next chunk break should go based on improvement in retrieving answers for downstream Question Answering.

Train the Chunker End-to-End on Question Answering

The goal is to train the chunker to directly optimize the retrieval process. To do this we use a preference optimization approach where the model has a choice between two sampled locations to set as the chunk divider, and we empirically determine which leads to better question answering and context retrieval.

Once the chunker is trained, it can be used as follows. Given a window of an unchunked portion of a document, the chunker predicts where the next chunk break should go. This creates a new document chunk and then the window slides forward. This process is repeated until the chunker reaches the end of the document.

We use Google’s “Flan-T5-Base” as our base model.

Data Collection and Processing

We use the SQuAD (Stanford Question Answering Dataset) for training and evaluation. SQuAD is a contextual question answering dataset. Each question-answer pair is mapped to a specific context text derived from Wikipedia that is needed to answer the question. Moreover, each QA pair also contains the index of where the answer comes from. We use this information in our training process.

To turn the dataset from a contextual question answering task, to a RAG task, we incorporate additional retrieval documents from Cohere/wikipedia-22–12-en-embeddings. We include the cohere documents and embeddings from texts derived from articles matching document titles in SQuAD. This gives us a corpus of 35,211 embedded document chunks from 401 wikipedia articles.

Model Configuration and Training Procedure

The training procedure goes as follows. For a given Question Answering pair and context we run the chunker greedily over the context document until we get a chunk containing the answer sentence. Once that happens we sample an alternative chunk location as well. In parallel, We continue to greedily chunk the context document from each of these two locations. This leaves us with two alternative chunkings of the context document with the first difference occurring near to the answer location.

We then embed the chunks from each of the two sample chunkings and do a standard retrieval process for answering the question. We use a subset of embeddings from Cohere/wikipedia-22–12-en-embeddings (link) as distraction documents for the retrieval process. This is similar to the actual RAG process wherein retrieval is done over an entire corpus.

We then assess the quality of the retrievals according to the metric tokens-to-answer-chunk. The retrieval process produces a ranking of relevant chunks. This metric sums the cumulative tokens in chunks retrieved up to and including the chunk containing the answer. This metric favors when the answer chunk is ranked higher as there will be fewer chunks, and thus tokens, preceding the answer. However, it also favors smaller chunks, all else being equal, as that means fewer tokens as well. Fewer tokens means RAG is more likely to get the answer in a model’s context window when actually tested.

This process is used to generate preference data for the chunker. The sampling process with the lower resulting Tokens-to-Answer is considered the preferred sample. We then train with Cross-Entropy loss.

Backpropagation and Optimizing the Loss

Once a batch of preference pairs is established, we re-run the forward pass for the chunking step in which the samples diverge. We then use a loss to optimize the preferred chunk location over the alternative one.

Initial Results

With the integration of weights and biases, we monitored average tokens for both the samples, Loss, and entropy values.

Observations

We see some hints that this approach could work. We can reduce the tokens-to-answer by 100+ tokens. However, this is a relatively small scale improvement from small scale experiments.

We see that learning can happen but is heavily limited. One challenge is that the training process is incredibly noisy. As we only train on the ability to answer a single question from the passage at a time, rather than all questions for a passage, the signal is more limited.

Limitations

Learnable chunking approaches could be an important part of the RAG process, however, this current implementation has a number of limitations.

1. A significant amount of time is required to perform Learned Chunking before passing the data to RAG, which may not be suitable for enterprise grade RAG applications or applications where the retrieval corpus is changing rapidly.
2. The section to chunk needs to fit in the encoder model’s context window. In this case, we use a T5 model with a 512 token context window. As LLM context windows get longer, we would expect the optimal size of chunked documents to increase as well.
3. The results are not indicative of a strong improvement from this process.
4. We have not tested on downstream QA.

Conclusion

Our Learned Chunking approach could be a powerful way to optimize a currently heuristic-based part of the RAG process. While the current version is not ready for use, future iterations could enable a Learned chunking approach.

For ideas, thoughts, and collaborations reach out to Krupa or Elan

Code: github.com/krupagaliya/approach_to_learned_chunking