Background

This post examines improvements made to ColiVara, our ColPali-based retrieval API. We focus on hybrid search and hierarchical clustering token pooling. By benchmarking these two approaches, we aim to evaluate their impact on latency and performance.

The conventional approach to handling documents for RAG or data extraction typically involves a multi-stage process: Optical Character Recognition (OCR) to extract text, Layout Recognition to understand the document structure, Figure Captioning to interpret images, Chunking to segment the text, and finally, Embedding to represent each segment in a vector space.

This pipeline is not only complex and computationally demanding but also prone to error propagation. Inaccuracies in any stage, for example, OCR errors or misinterpretation of visual layouts, can significantly degrade the quality of downstream retrieval and generation.

A more streamlined approach, as pioneered in the ColPali paper, leverages the power of vision models. Instead of complex pre-processing, this method directly embeds entire document pages as images, simplifying the retrieval process to a similarity search on these visual embeddings.

This approach eliminates the need for laborious pre-processing steps and potentially captures richer contextual information from the visual layout and style of the document. This innovative approach forms the foundation of our work in ColiVara.

This post details our research on using hierarchical clustering token pooling, building upon and extending the core ideas presented in ColPali.

Research Question

In our previous post, we looked into whether we could speed up our inference by using different similarity calculations.

We found that the bottleneck is directly linked to the number of embeddings for each document. The more embeddings there are, the longer the calculation takes. This relationship has a quadratic complexity of O(n*m), where m is the number of document vectors.

The question we wanted to answer was: Can we maintain the same state-of-the-art performance with either fewer document candidates or a significantly lower embedding count?

The goal was to see if we could improve latency by reducing the number of documents (using hybrid search) or the number of document vectors (using token pooling) without losing any performance.

Our metric, consistent with the ColPali paper, was the NCDG@5 score in ArxivQA and our typical API request latency using:

• Hybrid Keyword Search with Postgres native search capabilities to reduce the number of candidates for inference

• Token pooling at indexing time to reduce the total number of embeddings by averaging within similarity clusters

The ColPali paper uses late-interaction style computations to determine the relevancy between a query and documents. Here is a simple GitHub gist explaining how late-interaction style computations work: https://gist.github.com/Jonathan-Adly/cb905cc3958b973d7b3b6f25d9915c39

The key point is that late-interactions rely on multi-vector representation. Multi-vectors introduce significant storage and memory overhead and are computationally demanding to search through.

Baseline implementation

To get a realistic picture on real-life performance, we set the parameters as close to our production setup at ColiVara as possible. Embeddings were stored in Postgres with a pgVector extensions. Everything ran in an AWS r6g.xlarge (4 core CPU, 32g ram) and called from our python backend code hosted in a VPS.

We use ColQwen2 as the base model and colpali-engine (v0.3.4) to generate embeddings. It improves upon the base implementation of the paper with ~25% less embeddings and better performance. Our work builds on top of those improvements to enhance them further.

Results:

On ArxivQA using NCDG@5 and end to end latency we had the following:

• Average NDCG@5 score: 0.88

• Average latency: 3.58 seconds

The dataset is composed of 500 pages. This score matches the leader of the Vidore leaderboard and considered state of the art retrieval.

Hybrid search

Having run a few large RAG applications before, we were well aware of the power of hybrid search. The details of the implementation can vary, but the main idea is to quickly narrow down your candidate documents using Postgres search capabilities, then re-rank them with more advanced semantic search techniques.

In our implementation, we used gemini-flash-8B to create captions and add keywords to each document during indexing. At query time, we used the same LLM to convert the query into keywords, then employed standard Postgres search with a GIN index to retrieve the top 25 documents.

Results:

On ArxivQA using NCDG@5 and end to end latency we had the following:

• Latency: 2.65

• Score: 0.68

The good news is that we reduced our latency by about 25%. However, our performance decreased significantly. We believe that by adjusting all the parameters of the hybrid search implementation, we can improve performance. However, it will involve a trade-off between latency and performance. For instance, using a larger model like Qwen2-VL 72B for keyword extraction might enhance performance, but it will also be slower. Similarly, indexing the full text of the documents instead of just using captions or keywords might be better, but it would also slow things down.

Hierarchical Clustering Token Pooling

The next optimization we wanted to test is hierarchical clustering token pooling. Despite the complex name, it's actually quite simple. You take your embeddings and average them in clusters of similarity.

You start with a desired compression level, called the pooling factor. At a pooling factor of 1, you keep your embeddings as they are. With a factor of 2, you reduce your embeddings by half. At a pooling factor of 3, you save only a third of your original embeddings.

We got the idea from this excellent post by Answer.AI, which tried this approach using ColBert and found success. The optimal point seemed to be a pooling factor of 3, so we chose that.

The implementation to do this is ~10 lines of code. A standalone function that pools your embedding at index time or even after the fact.

def pool_embeddings(embeddings: torch.Tensor, pool_factor: int = 3) -> List[List[float]]:
    """
    Reduces number of embeddings by clustering similar ones together.

    Args:
        embeddings: Single image embeddings of shape (1038, 128)
                   Example with 4 vectors, 3 dimensions for simplicity:
                   [[1,0,1],
                    [1,0,1],
                    [0,1,0],
                    [0,1,0]]
    """
    # Step 1: Calculate similarity between all vectors
    # For our example above, this creates a 4x4 similarity matrix:
    # [[1.0  1.0  0.0  0.0],    # Token 1 compared to all tokens (same, same, different, different)
    #  [1.0  1.0  0.0  0.0],    # Token 2 compared to all tokens
    #  [0.0  0.0  1.0  1.0],    # Token 3 compared to all tokens
    #  [0.0  0.0  1.0  1.0]]    # Token 4 compared to all tokens
    # High values (1.0) mean tokens are very similar
    similarities = torch.mm(embeddings, embeddings.t())

    # Step 2: Convert to distances (1 - similarity)
    # For our example:
    # [[0.0  0.0  1.0  1.0],    # Now low values mean similar
    #  [0.0  0.0  1.0  1.0],    # 0.0 = identical
    #  [1.0  1.0  0.0  0.0],    # 1.0 = completely different
    #  [1.0  1.0  0.0  0.0]]
    distances = 1 - similarities.cpu().numpy()

    # Step 3: Calculate target number of clusters
    # For our example with pool_factor=2:
    # 4 tokens → 2 clusters
    target_clusters = max(embeddings.shape[0] // pool_factor, 1)

    # Step 4: Perform hierarchical clustering
    # This groups similar tokens together
    # For our example, cluster_labels would be:
    # [1, 1, 2, 2]  # Tokens 1&2 in cluster 1, Tokens 3&4 in cluster 2
    clusters = linkage(distances, method="ward")
    cluster_labels = fcluster(clusters, t=target_clusters, criterion="maxclust")

    # Step 5: Average embeddings within each cluster
    # For our example:
    # Cluster 1 average = [1,0,1] and [1,0,1] → [1,0,1]
    # Cluster 2 average = [0,1,0] and [0,1,0] → [0,1,0]
    # Final result: [[1,0,1], [0,1,0]]
    pooled = []
    for cluster_id in range(1, target_clusters + 1):
        mask = cluster_labels == cluster_id
        cluster_embeddings = embeddings[mask]
        cluster_mean = cluster_embeddings.mean(dim=0)
        pooled.append(cluster_mean.tolist())

    return pooled

This method of pooling is quite elegant and brilliant. Where - nothing else changes except the number of embeddings being generated.

Results:

On ArxivQA using NCDG@5 and end to end latency we had the following:

• Latency: 2.13

• Score: 0.87

This was as close to a free lunch as you will ever get. The storage cost went down by 66%, latency improved by about 40%, and there was very little performance loss. It was magnificent and beautiful, the kind of optimization that is usually only theoretical.

We decided to implement it in production and ran the entire evaluation suite. The final results were even more impressive, with up to 70% better latency on larger document collections and very minimal loss. You can see the full results here.

Conclusion

We tested and experimented with hybrid search and token pooling to improve our API's latency without sacrificing performance. Hybrid search improved latency by narrowing down results with keywords, but it reduced performance. On the other hand, hierarchical clustering token pooling greatly improved storage efficiency and latency with minimal performance loss.

Retrieval Augmented Generation (RAG) empowers large language models (LLMs) by integrating private documents and proprietary knowledge, unlocking their potential for nuanced and informed responses. However, efficiently extracting information from unstructured documents—especially those with complex visual layouts—remains a significant challenge.

The conventional approach to handling visually dense documents typically involves a multi-stage process: Optical Character Recognition (OCR) to extract text, Layout Recognition to understand the document structure, Figure Captioning to interpret images, Chunking to segment the text, and finally, Embedding to represent each segment in a vector space. This pipeline is not only complex and computationally demanding but also prone to error propagation. Inaccuracies in any stage, for example, OCR errors or misinterpretation of visual layouts, can significantly degrade the quality of downstream retrieval and generation.

A more streamlined approach, as pioneered in the ColPali paper, leverages the power of vision models. Instead of complex pre-processing, this method directly embeds entire document pages as images, simplifying the retrieval process to a similarity search on these visual embeddings. This approach eliminates the need for laborious pre-processing steps and potentially captures richer contextual information from the visual layout and style of the document. This innovative approach forms the foundation of our work in ColiVara.

This post details our research on optimizing and benchmarking various similarity calculation methods, building upon and extending the core ideas presented in ColPali.

We store vector embeddings in a PostgreSQL database, employing a one-to-many relationship between vectors and their corresponding pages.

Research Question

The ColPali paper uses the late-interaction style computations to calculate the relevancy between a query and documents. The question we wanted to answer, could we get the same results using other similarity computations with better latency?

Our metric, to stay consistent with the paper was the NCDG@5 score in DocVQA and our typical API request latency using:

• Late-interaction Cosine Similarity

• Hamming distance with Binary Quantization of vectors

• Hamming distance as above with late-interaction re-ranking

The core idea in simple terms in the ColPali paper implementation of late-interaction is:

• We have query embeddings: an array with n vectors, each of size 128 (floats).

• We have document embeddings: an array with 1038 vectors, each of size 128 (floats).

• For each query vector (n), we find the most similar document vector by computing the dot product. Here is simple code for straight-forward dot product similarity.

    import numpy as np
  
    # Define the query vector
    query = np.array([1, 2, 3])
  
    # Define the document vectors
    document_a = np.array([4, 5, 6])
    document_b = np.array([7, 8, 0])
  
    # Compute the dot product between the query and each document
    dot_product_a = np.dot(query, document_a)  # 1*4 + 2*5 + 3*6 = 32
    dot_product_b = np.dot(query, document_b)  # 1*7 + 2*8 + 3*0 = 23
  
    # Output the results
    print(f"Dot product between query and document a: {dot_product_a}")
    print(f"Dot product between query and document b: {dot_product_b}")
    # document a is more similar than document b to our query
  

• Late interaction works like this: take each query vector, find which document vector it matches best with (using dot product), and add up all these best matches. Here's a simple example that shows exactly how this works:

import numpy as np

# One query generates multiple vectors (each vector is 128 floats)
query_vector_1 = np.array([1, 2, 3])  # simplified from 128 floats
query_vector_2 = np.array([0, 1, 1])  # simplified from 128 floats
query = [query_vector_1, query_vector_2]  # one query -> n vectors

# One document generates multiple vectors (each vector is 128 floats)
document_vector_1 = np.array([4, 5, 6])  # simplified from 128 floats
document_vector_2 = np.array([7, 8, 0])  # simplified from 128 floats
document_vector_3 = np.array([1, 1, 1])  # simplified from 128 floats
document = [document_vector_1, document_vector_2, document_vector_3]  # one document -> 1038 vectors

# For each vector in our query, find its maximum dot product with ANY vector in the document
# Then sum these maximums

# First query vector against ALL document vectors
dot_products_vector1 = [
    np.dot(query_vector_1, document_vector_1),  # 1*4 + 2*5 + 3*6 = 32
    np.dot(query_vector_1, document_vector_2),  # 1*7 + 2*8 + 3*0 = 23
    np.dot(query_vector_1, document_vector_3)   # 1*1 + 2*1 + 3*1 = 6
]
max_similarity_vector1 = max(dot_products_vector1)  # 32

# Second query vector against ALL document vectors
dot_products_vector2 = [
    np.dot(query_vector_2, document_vector_1),  # 0*4 + 1*5 + 1*6 = 11
    np.dot(query_vector_2, document_vector_2),  # 0*7 + 1*8 + 1*0 = 8
    np.dot(query_vector_2, document_vector_3)   # 0*1 + 1*1 + 1*1 = 2
]
max_similarity_vector2 = max(dot_products_vector2)  # 11

# Final similarity score is the sum of maximum similarities
final_score = max_similarity_vector1 + max_similarity_vector2  # 32 + 11 = 43

print(f"Final similarity score: {final_score}")

Now - as you can imagine, this is computationally heavy at scale (or so we thought!). So, we looked for ways to see if we can make it faster and maybe even better.

Baseline implementation

To get a realistic picture on real-life performance, we set the parameters as close to our production setup at ColiVara as possible. Embeddings were stored in Postgres with a pgVector extensions. Everything ran in an AWS r6g.xlarge (4 core CPU, 32g ram) and called from our python backend code hosted in a VPS.

We implemented the paper late-interaction calculation as a Postgres function as such:

 CREATE OR replace FUNCTION max_sim(document halfvec[],
                                   query halfvec[]) returns DOUBLE PRECISION
AS
  $$ WITH queries AS
  (
         SELECT row_number() OVER () AS query_number,
                *
         FROM   (
                       SELECT unnest(query) AS query) ), documents AS
  (
         SELECT unnest(document) AS document ), similarities AS
  (
             SELECT     query_number,
                        (document <#> query) * -1 AS similarity
             FROM       queries
             CROSS JOIN documents ), max_similarities AS
  (
           SELECT   max(similarity) AS max_similarity
           FROM     similarities
           GROUP BY query_number )SELECT Sum(max_similarity)
  FROM   max_similarities;

The magic line here is (document <#> query) * -1. The <#> symbol is from pgVector and we multiply with -1 as gives negative inner product by default. We also use halfvecs to store embeddings for efficiency.

Results:

On DocVQA using NCDG@5 and end to end latency we had the following:

• Average NDCG@5 score: 0.55

• Average latency: 3.1 seconds

The dataset is composed of 500 pages.

Late-Interaction Cosine Similarity

Having ran a few large RAG applications before, we were intimately familiar with how cosine similarity works. Cosine similarity normalizes for vector length and gives a few nice additions.

• Values are always between 0 and 1

• This means vectors pointing in the same direction but with different magnitudes will be considered similar

• It is more computationally more expensive than dot product as it needs to calculate vector magnitudes (requires square roots) but also more performant

Our first attempt was meant to establish what would happen if we kept everything the same, but substituted dot product with cosine similarity.

So, we modified our function by replacing the dot product with cosine similarity.

Before:

(document <#> query) * -1 AS similarity

After:

1 - (document <=> query) AS similarity

Then - we ran our evaluations again.

Results:

• Average NDCG@5 score: 0.55

• Average latency: 3.25 seconds

So - our latency went up, but our scores didn’t really improve. It was good to see though that the main latency-driver is not the math behind the calculation. We could use it, and will likely be okay. However, both the latency benefits and the improvements are marginal.

Binary Quantization and Hamming Distance

Binary quantization is a compression technique for vector databases. It converts 32-bit floating point numbers into binary values (1 bit), achieving significant memory savings. The conversion process typically preserves the relative relationships between vectors while sacrificing some precision. For example:

embeddings = [0.123, -0.456, 0.789, -0.012]
quantized_embeddings = [1, 0, 1, 0]

To measure similarity when we have bits as our embeddings, we use the hamming distance. Hamming distance is a measure of the number of differing bits between two binary strings of equal length. It's calculated by comparing the corresponding bits of the two strings and counting the positions at which the bits differ.

First - we converted all our embeddings to be binary and stored them in their own column. Then - we changed our similarity function to be as the pgVector-python documentation recommends.

1 - ((document <~> query) / bit_length(query)

Finally - we ran our evaluations.

Results:

Average NDCG@5 score: 0.54

Average latency: 3.25 seconds

The good news is that the precision loss is minimal. We are hesitant to embrace it fully even with these results, as the loss can be unpredictable. This is essentially a measure of bit diversity, which partly depends on the data being embedded—so your mileage may vary.

The interesting tidbit here is latency is slower, as slow as Cosine similarity. We can almost conclude that if we want to make the process faster, what similarity calculation we should use is probably irrelevant.

Binary Quantization with Re-ranking w/ 100 documents

The final evaluation we wanted to conduct before wrapping up this experiment was to see if we could use Hamming distance and then re-rank using dot product late interactions. If our theory is correct—that the type of similarity calculation doesn't matter much—we should expect a slight increase in latency (since we're doing two calculations) and no significant improvement in performance.

That's exactly what happened.

Results:

• Average NDCG@5 score: 0.55

• Average latency: 3.5 seconds

The bottleneck

Postgres is highly optimized, so adding a square root or doing a bit of extra multiplication has a minimal impact. It's an obvious realization in hindsight, but without conducting evaluations and this experiment, we wouldn't have been able to prove it.The real bottleneck is here:

FROM queries CROSS JOIN documents

the CROSS JOIN is the major bottleneck ! Let's break down the scale:

For each comparison:

• Query: n vectors

• Document: 1038 vectors per page

• Corpus: 500 pages

So for ONE query against ONE page we're doing: n * 1038 vector comparisons (CROSS JOIN)

And for the full 500 pages corpus: n * 1038 * 500 vector comparisons

This creates a massive cartesian product. For example, if n=10 (for query):

• 10,380 comparisons per page

• 5,190,000 comparisons for 500 pages

The CROSS JOIN is creating this quadratic complexity O(n*m) where:

• n = number of query vectors

• m = number of document vectors

So - the math behind the comparisons is relatively inconsequential. The real optimization has to happen before we reach this point.

What can be done?

The crucial optimization insight that we got out of this experiment, is to hyper-focus on the number of pages/documents we are running through our similarity function. Nothing else really matters.

if instead of CROSS JOIN documents (1038 vectors * 500 pages), we first filter:

CROSS JOIN (
    SELECT * FROM documents 
    WHERE page_id IN (
        -- something to pre-filter to maybe 10-20 relevant pages
        -- instead of all 500
    )
)

All of the sudden, we gain massive performance gains. The insight here is for this specific problem: 1038 * 10 pages is so much better than 1038 * 500 pages.

It's like finding a book in a library:

• First narrow to the right section (cheap and quick search)

• Then look for the exact content (expensive vector similarity)

We currently use and recommend advanced filters in ColiVara to reduce the document set. When a user uploads a document, they can add arbitrary JSONs as metadata. It can be very complex or simple. At query time, they can use filtering to reduce the corpus size. This is a powerful technique, but can be complex to setup.

Future work

There is a couple of approaches we planning to tackle to reduce the number of documents without any effort from the user side. The first is a few experiments using Postgres full text search. This will take us away from vision models, but could be straightforward win.

The second is to use vision LLMs generated metadata of documents to store semantic information about the documents automatically at indexing. We can then use existing filtering to reduce the corpus size via that metadata.

The main problem vision models solve in our experience, is that traditional OCR pipelines can’t capture the full visual cues in documents. Here, we don’t want to capture everything, just enough to narrow the search and then vision models will do the rest.

Conclusion

We benchmarked and experimented with several similarity calculations setups for our RAG API using vision models. The differences are marginal and doesn’t meaningfully affect performance or latency.

Due to the large size of embeddings produced, the best performance optimization is to reduce the number of documents that go through the similarity calculation, rather than optimize that similarity calculation itself.