Late Chunking for Enhanced Retrieval in Long-Context Embedding Models

[00:00:00]in a previous video we looked at contextual retrieval from anthropic which is their context enhancement technique for improving rag systems but there's 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 embedding 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 this huge Max token for some of the newest embedding models but there is one major issue with these in a standard rack pip pipeline irrespec 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 um B L from Wikipedia if you use a sentence level uh 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's in 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 and embedding process in the standard uh rack pipeline you convert your uh text into smaller chunks then pass each through a 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 uh 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 the other chunks in the same document the late chunking approach reverses this instead of dividing your document into chunks and then Computing in wedings 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 uh this is called late chunking and it's directly related to another approach called late interaction which is this cber 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 viit 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 trunking approach which is about 5 GB but if you were to do uh uh late interaction or cobbert based um multiv Vector representation you will need about 2.5 terabyt which is pretty huge the reason is that you're storing uh these embeddings for each token individually but this late chunking approach which gives you the best of both worlds not only it preserves the context in your final uh 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 embeding for your whole document and then do chunking at a later step if you have a huge document you want uh an embedding model that uh has a long context window or it has um a large number of tokens that it supports as an input and that's where this long context embeding 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 3 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 inje 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 uh chunking with their embedding model and the EM the choice of embedding model plays a critical role which we are going to uh 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 uh close to 30% boost over the base 9 now the uh main requirement for this is that you need to have a lot of text in your uh input to the embedding model after that the uh 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're 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 uh the contextual retrieval comparison in their block post so according to the blog post the anthropics method is Brute Force approach to address the issue of lost context here is 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 uh the full document then the llm adds the relevant context to each chunk and then this U results in richer and more informative embeddings essentially this is just another uh chunking technique which improves the context for each of the chunk but as you can see since you're sending the whole document this is going to be very expensive uh in terms of the cost time that it needs as well as the storage now anthropic recommends to use um 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 model 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 pulling of the final embeddings and we are talking about 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 advanc 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 embeding models there's a second part to it which is what late chunking really is and what it's not part two then there is another blog post from the vivit team which is a vector store late chunking balancing precision and cost in Long context retrieval we looked at some of the images from this blog post all three block 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 uh worth the read uh so I'm going to put this in the video description they also have a GitHub repo late chunking of short chunks and long context Elling 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 uh 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 segmented API which uses regular Expressions to chunk your documents you can use this sentence level embeddings sentence level chunking or the recursive character uh chunk Splitter from a length chain 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 is the text that we want to chunk this is an entry from Wikipedia regarding Berlin this 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 embedded and then you do the sentence that will chunking so this single function will uh 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 uh is still based on the uh chunks you created even from the late chunking process you can uh 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 of 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 82 for the traditional chunking approach the similarity drops to7 similarly the third sentence which uh refers to Berlin by this city the similarity is pretty high for late chunking but it's uh substantially low for the traditional chunking approach now this is a very simple example but you can uh 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 uh have proposed now one other thing I wanted to mention is that late uh chunking is bidirectional if certain information is before a specific chunk it can only preserve that information but since it's bidirectional 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 B directional and I think it makes it even more powerful so this shows that long context models are important both for LMS 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 100,000 or more but you also want to look at the memory usage for example um this multilingual gamma 2 embedding model needs about 34 GB 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 geni embeddings might be a potential option now a potential approach would be that you take something like this long context approach along with the uh contextual retrieval approach proposed by anthropic and combine them together in your retrieval pipeline that way you can get best of the uh Both Worlds anyways um 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 agentic systems and what things to consider when You're Building any llm based applications if that's something which interests you I'd highly recommend to 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