ColPali: Vision Language Models for Efficient Document Retrieval

There's a new paper out called Kalpany Efficient Document Retrieval with Vision Language Models. And they propose a very interesting approach which can potentially change the way we do retrieval in RAG today. They propose to use vision models for retrieval and apart from excellent retrieval results, it also gives you explainability and the results are pretty impressive compared to the existing approaches that we have. So in this video I'm going to walk you through some of the current challenges with RAC systems. What is the potential approach that is being presented here? We're going to look at some results and then we will get to play around with this new model as well. So first let's understand some of the issues with the existing RAG systems. One of the aspects of RAG is parsing the data. Most of the data that we deal with is in either PDFs or HTMLs or CSV and it's actually extremely hard to parse data out of something like a PDF. And this is a part of the Rank pipeline that doesn't get a lot of attention. But there are special libraries like Llama Parse or Unstructured IO that can help you parse information from PDF files. However, it is actually a cumbersome process. So for example, the typical process is going to look something like this. You have your PDF files, then you are going to run an OCR model for optical character recognition. Then you also want to have a layout detection model because these PDFs can have very weird structures in there. And after that, once you parse the text from the PDF then you will need to run it through a chunking process. So you need to decide what chunk size to use, what type of chunking strategy to use. After that you're going to embed those chunks and put them in a vector store. So there are a lot of different moving parts to this pipeline and the errors from each of the part can accumulate over time. Instead of doing this multi part approach, they propose a relatively simpler solution. So they propose to take the images of the pages in a PDF file, then use a vision model to embed those using a Vision encoder, and then use a model like polygamma, which is a vision language model to extract information. So essentially you are doing a retrieval directly on the images that were created from the PDF files. So how good is this approach? Well, they had to create a new benchmark because there is currently no way to compare vision models and dense embedding models for retrieval. So they created a new benchmark that has a number of different PDF files with information related to different fields and it outperforms all the existing methods. So the existing methods that they compared was to use something like unstructured I O, which is a pretty amazing library for parsing PDF files, extract text on that, and then they compare it with keyword based approaches. So BM25 is kind of state of the art when it comes to keyword based retrieval. And they also compared it against dense embedding based approaches. So BGE M3 is one of the best multilingual multitask embedding model for retrieval purposes. They also compare the same thing with unstructured IO extraction plus doing OCR and captioning on top of it. And now since they created this new data set, they had the ability to compare it directly with other vision models and vision embeddings. And the results, well, it shows it outperforms every other approach by a huge margin. There's not even anything close to it. And we are going to talk about how it works in a lot more detail in a second. But there's something else I want to highlight which is very important for anybody who is building a RAG system, and that is if you look at text based approaches, BM2025 seems to be as good or even better in retrieval compared to dense embedding based approaches. That means a traditional keyword based approach can be as good or maybe even outperform your embedding based retrieval for certain applications. So if you're building a RAC system, you definitely want to make sure that you include keyword based mechanism along with your embedding base retrieval. Now, I have talked about the issues with embedding based retrieval in a number of my videos, and in one of them I highlighted a potential alternative, which is Colbert, which looks at late interactions of the tokens. So it basically creates a multi vector representation for a given chunk rather than using a single dense embedding vector. And in fact, the CAW Poly paper uses that approach to improve the performance of vision based retrieval. Okay, so here's the architecture of Call Poly, which is based on Polygamma 3 billion model from Google. And we're going to go into details of different parts of this architecture, but the main component is that each image is a page in a PDF file. The first step is to process these images. The process of creating patches involves dividing the input image, which is basically a page of a PDF file, into a grid of smaller equal size sections. For Call Poly, since it's using Polygamma 3 billion, the image is divided into a grid of 32 by 32 and that means for each image we are going to have 1024 different patches. Next, for each patch, we are going to create embeddings. And here is how the embedding creation process looks like. So first, each patch is embedded into a higher dimensional vector space using a linear projection. This step helps to capture initial features from the raw pixel data. Next, those embeddings or patch embeddings are then processed by the Vision transformer which applies self attention mechanism to capture the relationship between patches. So this step allows the model to capture relationship between different parts of the image, which means it will provide context to each of the patch that it's seeing in the image. Now in the next step, the output of the Vision transformer is fed into the language model, which is polygamma in this case that further processes the visual information. It potentially aligns with the textual representation as well. So this actually helps the model to learn not only the representation of text that is present, but it can also look at plots, images and tables in the document or on that one page without actually needing to extract these different elements. Using some other parsers like Llama Parser Unstructured I O, the output of this Vision language model is projected into a lower dimensional space. In this case it uses a vector of 128 dimensions for each of the patches. Now here is how this whole process is going to look visually. And the process is applied to all 1024 patches of the input image. This will result in a set of 1024 embedding vectors where each has a dimension of 100 in units. This is basically a multi vector representation of each of the image or each of the page. Very similar to how Coldbot does multi vector representation of a texture chunk compared to like a dense embedding vector. So the key advantage of this approach is that it captures both the local features so from individual patches and global context through the Vision transformer and language model processing. This allow the call poly to understand complex visual layouts, texts, images within the document without actually relying on any previous pre processing and enabling it to be a more effective retrieval for visually rich documents. Okay, so we talked about how we do this multi vector representation for a single page. But if the document contains 10 pages, we're going to do the same process for each of the page. Now how does the retrieval looks like and where does the late interaction takes place? So in this section I'm going to walk you through that process. So for input query we're going to first tokenize the input query and each token is encoded into 128 dimension vector using the same polygonal model, because that's a vision language model. So this results in a set of query tokens or embeddings. Here we have three embeddings for this example use case. Then we take our image and we look at the multi vector representation. So just to show you an example, instead of 1024 patches, we're just looking at about 25 patches here. And each of these patches have associated 128 dimension vector. So there the size of the vector representation, both for the query as well as for each of the patches is similar. Then we do a similarity matrix. So similarity computation. And the way it works is that for each query token, its similarity is computed with all the document patch embeddings. This creates a similarity matrix where each cell represents similarity between a query token and a document patch. The approach is very similar to how Colbert works. And the similarity is typically computed using a dot product or cosine similarity. You can use other metrics as well, but those are the two that are usually used. So similar to the later interaction in Colbert, we do max pooling. And the idea is for each query token we keep only the maximum similarity score across all patches. So this helps us identify which part of the document is most relevant to each query. And then we do a final aggregation step. So the max pooled similarity for each query token are summed to get the final similarity score between the query and the document. Now we're going to do this for each of the pages that are present in the document. And this will let you retrieve the most relevant pages. So we're not looking for specific paragraphs, but we are going to be able to retrieve specific pages in which the information that we are looking for is contained. Okay, so here's how it's going to look like. You are going to have your query and then using Polygamma, it will retrieve the most relevant pages. So out of 10 pages page document, it will just pick, let's say the top three pages in the ranked order. Then you can do a couple of things. Either you can do text retrieval on those pages and take this as a context and provide it into LLM to generate responses. That's like one way, but I think it's not going to be a very efficient way. The second way would be to basically feed those pages as images into vision language model or a multimodal model, such as in Gemini Flash or GPT4O or even something like Sonnet 3.5. Now you can potentially Take images of all the pages that are present in your database and feed these directly. However, that is going to be extremely expensive. Imagine that you have thousands of pages in a document or a set of document. But this can help you reduce the number of documents that you are using. Feeding for final generation. So it can substantially reduce the number of documents this multimodal model has to process. And that is going to be the generation part. The paper only talks about the retrieval part, which is basically how do you pick the top three or five pages out of all the pages that are in your corpus? An obvious question is going to be does this processing add to the latency of the system? Now it turns out that it's very efficient when it comes to the indexing process which is basically creating this multi vector representation for images of the PDF files. So they did a comparison. If you are going to do OCR plus layout detection along with the normal chunking and embedding creation process, that takes about 7.22 seconds per page. For col poly this takes about only 0.40 seconds or 39 seconds per page, which is extremely efficient when it comes to indexing process. However, it is less performant in terms of the compute time it takes when it comes to the query part. So for a dense embedding vector based retrieval it takes about 22 milliseconds per query. But here it adds 8 seconds. So it's going to take about 30 milliseconds per query. The model is available on Hugging Face and here's a Python script of how to use it. If there is interest. I'm going to create another video in which we're going to actually implement this and see how we can run this on our own data set. There's a very helpful guide from a Vispa blog. Vispa is another Vectra store provider. They have their own cloud and actually there is a Google Colab notebook if you are interested. I'm going to put a link to this. If you want to test it out yourself, you can do that using this Hugging face spaces. Now depending on the GPU availability it might take some time, but essentially you can upload your own document if you want and then index them. They have also provided an example document which is Climate Youth magazine. So we're going to click on that. You can see that here's the file that it's going to be using. Then you're going to click on index documents. This will create this multi vector representation we talked about for each page and for the whole document. Now the process can take some time depending on whether a GPU is available or not. Right now it's running on a free gpu, so this could take some time or timeout. It actually indexed everything. The document contains 32 pages and we're going to run the example query that they have provided here. So how much tropical forest is cut annually and we are asking it to retrieve, let's say the top three pages that I think are more relevant to this query. So you can click on search query that will do the retrieval process. So basically using those late interaction it will try to find the most relevant pages in our PDF. Alright, so it gave us three different pages. The first one talks about an area of tropical forest about the size of Greece is cut down every year. And so it does talk about the rainforest in this page. The second page is related to soil are home to over a quarter of all living species. Right. And if we look at this page, I think it has specific information about rainforests. I'm going to be running some more tests. But an added advantage of using this vision language model is explainability. So in the paper they have this image since it's using a vision model and the vision transformer can pay attention to different patches. So you can actually see which patch is more focused for a specific query. So here is an example. Which hour of the day had the highest overall electricity generation in 2019? So you can see that most of the attention is first on the hours. So basically the model is actually specifically looking for different hours and also focusing on these average hourly generation and numbers as well. And then since the query talks about 2019 so there's more attention to 2019 and hour because the query is asking about a specific hour of the day. And then next I think some more attention on the plot itself. Right. So you can see that not only it can do a really efficient retrieval but but it also gives you explainability of which specific patch the document or the model is looking at in the document. So extremely helpful approach. The next step would be to connect this to a multimodal model to do the actual generation. Very interesting approach. I'm planning on creating more content on it and showing you how to run this locally. Because it's just a 3 billion parameter model so it's going to be very cost efficient if you want to run it on your local machine. Let me know if you are interested in that. Make sure to subscribe to the channel so you don't miss out on that video. I hope you found this video useful. Thanks for watching. And as always. See you in the next one.