Stop Losing Context! How Late Chunking Can Enhance Your Retrieval Systems

In a previous video we looked at contextual retrieval from Anthropic, which is their context enhancement technique for improving RAG systems. But there is another technique called late chunking in long context embedding models which I think is a lot more interesting and can be significant. Embeddings is one of the most critical component of any retrieval system, but are often ignored or misused. When you are selecting an embedded model, you need to consider two very important parameters. One is the max tokens, which is basically the context window. And the second one is the embedding dimension which is the output size of the embedding vector. You have probably seen these huge max token for some of the newest embedding models, but there is one major issue with these. In a standard RAC pipeline, irrespective of the size of the chunk, the output size of your embedding vector is going to remain the same. So whether you are embedding five tokens or 5,000 tokens, the output is going to be exactly the same. That means these embedding models are going to be compressing a lot of information. For long input chunks, in most cases you want to use smaller chunks, but that has its own problems. To understand that, let's look at this example. Here is a small paragraph regarding Berlin from Wikipedia. If you use a sentence level chunking strategy, essentially each sentence is going to become its own different chunk. You can see that if you embed these different chunks separately, you will lose the contextual information. So it and the city is referring to Berlin. But if this chunk is in isolation, it's going to lose that context. The contextual retrieval approach tries to summarize the documents and add contextual information to each chunk. But there is a better approach and that is late chunking. In order to understand this, let's look at normal chunking in embedding process. In the standard RAC pipeline, you convert your text into smaller chunks, then pass each through neural network which is a transformer model. The output is going to be embeddings for each of the tokens individually. So you decompose your chunks into tokens and compute embeddings for each of the tokens. There is contextual information in these tokens, but they are limited to the chunk that it belongs to. After that, you compute mean pooling. You put all these token embeddings together, compute the mean and get the final output. Now this preserve the contextual information, but only within this chunk. It's independent of the other chunks in the same document. The late chunking approach reverses this instead of dividing your document into chunks and then computing embeddings, you first take the whole text of the document, pass it through a transformer model, and for each of the token you will get an embedding representation. But now, since you're passing on the whole document, these tokens have contextual information. And that's not only limited to a single chunk, but each of the token is going to contain contextual information of the whole document. After this step, you can do chunking. Essentially you create chunks of the original text, take the corresponding tokens, compute mean pooling for those, and get the final representation. Now, since we're doing the chunking process at a later stage, this is called late chunking and it's directly related to another approach called late interaction, which is this Colbert based approach which I have covered in my previous videos. This is probably the best approach for retrieval, but it comes with a cost of storage needs. In this case, in the final step, you don't do the pooling step, but you take individual token embeddings and store those token embeddings. Now here is a blog post from the VIVIT team. And in here they show that if you embed about 100,000 documents with the same number of embeddings, you will need about 1.6 million vectors for a naive chunking approach, which is about 5 gigabytes. But if you were to do late interaction or Colbert based multi vector representation, you will need about 2.5 terabytes, which is pretty huge. The reason is that you are storing these embeddings for each token individually. But this late chunking approach gives you the best of both worlds. Not only it preserves the context in your final chunking process, it also gives you about the same storage needs as the naive chunking approach. Now, later in the video, I'm going to show you how you can implement this in your own applications and start using it. But before that, let's understand the role of long context in these embedding models. Especially when it comes to late chunking. The name late comes because you compute embeddings for your whole document and then do chunking at a later step. If you have a huge document, you want an embedding model that has a long context window or it has a large number of tokens that it supports as an input. And that's where this long context embedding models become very important. So this approach was proposed by a company called Jenna AI and they also have their own embedding models. The latest one is Jenna Embeddings version three, which has a context window or max number of tokens that it can process about 8,000 tokens. 8,000 tokens equates to almost 10 pages of text. So that means you can embed about 10 pages of text and it will preserve the contextual information within the text. Now another related question is what about the chunk size? Whether you're doing chunking in the beginning or at the end of this ingestion pipeline. Well, they show some very interesting results. Now according to their results, it doesn't really matter how you chunk your documents. So they're comparing three different approaches. One is these late chunking with their embedding model and the embed. The choice of embedding model plays a critical role which we are going to address in later in the video. But then they compared this with sentence level chunking plus semantic chunking. But they showed that if you do sentence level chunking with late chunking approach plus the embedding model, that gives you close to 30% boost over the baseline. Now the main requirement for this is that you need to have a lot of text in your input to the embedding model. After that the nature of chunking process you use does not really matter. Now keep in mind, we will probably need a lot more validation from independent sources. But this is the result at least they are showing. And the results on these four different benchmarks shows that late chunking with their newest embedding model gives you the best possible or state of the art results at the moment. Now how does their approach compares to contextual retrieval? So I'm actually happy that they address the contextual retrieval comparison in their blog post. So according to the blog post, the Anthropics method is brute force approach to address the issue of lost context. Here's how it works. I have covered this in a couple of my videos if you are interested. I'm going to put links to those videos. But essentially each chunk is sent to the LLM along with the full document. Then the LLM adds the relevant context to each chunk and then this results in richer and more informative embeddings. Essentially this is just another chunking technique which improves the context for each of the chunk. But as you can see, since you are sending the whole document, this is going to be very expensive in terms of the cost time that it needs as well as the storage. Now Anthropic recommends to use their prompt caching which will reduce your cost by about 90%. And in another video I also show you how you can use the same approach with other models. Not only anthropic, but there are a couple of issues. First, this approach is still prone to the chunking boundaries or chunking strategy. The late chunking does not seem to be reliant too much on the chunking boundaries. The good news is that late chunking is not just limited to their own embedding model. You can essentially pick any embedding model that has long context and it has support for mean pooling of the final embeddings. And we are talking about token level embeddings now. All of these techniques are geared towards improving retrieval in retrieval augmented generation systems. If you want to learn more advanced techniques that you can use in rag, you can check out my course on advanced RAG techniques called RAG Beyond Basics. Information is in the video description. Before looking at the code or how you can implement this, I want to highlight these three blog posts. One is late chunking and long context embedding models. There's a second part to it which is what late chunking really is and what it's not. Part 2 Then there is another blog post from the VIVIT team which is a vector store late chunking, balancing precision and cost and long context retrieval. We looked at some of the images from this blog post. All three blog posts are worth reading and I highly recommend to check them out. And they even released a technical paper Late Chunking Contextual Chunk Embeddings using Long Context Embedding Models. It's worth the read so I'm going to put this in the video description. They also have a GitHub repo late chunking of short chunks in long context embedding models. So it talks about why late chunking is important. How it works. We looked at some of the images from here. Now what if you want to use this approach in your own applications or Rack pipelines, there is a very detailed notebook that you can use it for your own applications. Let me quickly walk you through this. First we need to install the transformers package. Then we are going to be loading their embedding model in this notebook. They're using the V2, but I think you can also use the V3 because that is available now we load both the tokenizer plus the actual model. There are two ways you can chunk your documents. Here's an example function that chunks your documents on sentence level. They have their own free segmenter API which uses regular expressions to chunk your documents. You can use this sentence level embeddings, sentence level chunking or the recursive character chunk Splitter from LangChain. If you decide to use the API, here's a function that will help you use their API endpoint. Again, it's free to use. They also have other free API endpoints, especially for web scraping, which is neat. Now here's the text that we want to chunk. This is an entry from Wikipedia regarding Berlin. They show you can chunk it by sentences. Here's how the chunks are going to look like if you do chunking, chunking or sentences. We basically have three different sentences out of the text based on this chunking by sentences function. After that they take those sentences and embed them using the traditional chunking process, which is essentially take this whole sentence and pass it to an embedding model to chunk it. Or you take the whole document first, embed it, and then you do sentence level chunking. So this single function will do both of the things for you. This function essentially does the late chunking for us. Now, if you directly call the encode function of the model on the chunks, you will get the traditional embedding representation of those individual chunks. And since the output is still based on the chunks you created even from the late chunking process, you can use traditional vector store to store those vectors. So there is no difference in terms of how this is going to look like or how you store these new embeddings. And you can store them in a traditional vector store. They are computing the similarity for Berlin as a single individual token to these new embeddings. If you look at the first sentence which talks about Berlin directly and mentions Berlin, both the traditional as well as late chunking outputs are going to give you very similar similarities. But then the second sentence is it's more than 3.85 million inhabitants. It is referring to Berlin, but it doesn't directly mentions it. In this case, if you look at the late chunking approach, the similarity is 0.82. For the traditional chunking approach, the similarity drops to 0.7. Similarly, the third sentence, which refers to Berlin by this city, the similarity is pretty high for late chunking, but it's substantially low for the traditional chunking approach. Now this is a very simple example, but you can think of if these chunks are much bigger, there's going to be a lot more loss of information in a naive chunking approach compared to the late chunking approach that they have proposed. Now one other thing I wanted to mention is that late chunking is bi directional. If certain information is before a specific chunk, if it can only preserve that information. But since it's bi directional and it's looking at the whole document when it's embedding information. Whether the information is before or after a specific chunk, that chunk will still preserve some of the information related to that chunk. So that's why it's going to be bidirectional and I think it makes it even more powerful. So this shows that long context models are important both for LLMs as well as for these embeddings. You actually want to experiment with these One more thing to highlight. You will see some of these models have a context window or max number of tokens of 32,000 and even some with hundred thousand or more. But you also want to look at the memory usage. For example, this multilingual Gamma2 embedding model needs about 34 gigabytes of VRAM, which is a lot more than a normal LLM. You need a dedicated A100 or H100 just to run these models. So that's why you want to have a compromise and the Jenna embeddings might be a potential option. Now a potential approach would be that you take something like this long context approach along with the contextual retrieval approach proposed by Anthropic and combine them together in your retrieval pipeline. That way you can get best of the both worlds. Anyways, if you found this video useful, make sure to subscribe to the channel. I create a lot of content geared towards developers and general populations, but the focus is understanding how to build better Rack systems, how to build better Agentix systems, and what things to consider when you are building any LLM based applications. If that's something which interests you, I highly recommend do subscribe to the channel. Also, if you want to support my work, there are a number of different links in the video description. Do check them out. I hope you found this video useful. Thanks for watching. See you in the next one.