Optimizing ColPali Retrieval at Scale

ColPali-style retrieval is one of those things that feels magical when it works: you embed a query, compare it against patch-level vectors from document pages, and get very strong visual-text retrieval quality.

The catch is cost.

In this project, each page produces about 1,030 token vectors. With late interaction (MaxSim), that means a lot of comparisons per query. It is fine on small demos, but once your corpus grows, latency and storage pressure quickly become real production issues.

If you are building RAG or enterprise search, this tradeoff shows up immediately:

• Product wants high recall.
• Users want fast responses.
• Infrastructure wants lower memory and compute load.

The goal of this work was simple: keep ColPali-level retrieval quality while making search faster and lighter.

I tested multiple optimization paths on a reproducible setup in Qdrant:

• dual-axis pooling (row + column),
• token clustering (KMeans centroids),
• MUVERA prefetch + rerank,
• and ablations across prefetch depth and reranker choices.

The short version: the initial intuition (dual-axis pooling) did not hold up, and token clustering became the best practical strategy.

1- Basline Setup

Before optimizing anything, I established a baseline that all methods are compared against.

Dataset and evaluation protocol

• Dataset: vidore/docvqa_test_subsampled (prepared via data/download_vidore.py)
• Working set in this experiment ( 500 page images, 500 query rows and one relevance label per query in this setup
• Metrics ( Recall@10, MRR, NDCG@10 andmean query latency (ms))

Retrieval baseline

The baseline is full multivector ColPali search (used VAGOsolutions/SauerkrautLM-ColQwen3–4b-v0.1 as the document/ query embedding model) in Qdrant with MaxSim scoring.

SauerkrautLM-ColQwen3–4b-v0.1 achieves 90.80 NDCG@5 on ViDoRe v1, making it the #2 overall among 128-dim models and the best in the Large (3–5B) category for ViDoRe v1.

Table : Full-Baseline Results

Collection/index setup (reproducible)

The experiment collection includes baseline and optimized vector slots:

• colpali_original (SauerkrautLM-ColQwen3–4b-v0.1)
• pooled variants (row_pooled, column_pooled, hierarchical_2x, hierarchical_4x)
• quantized variants (scalar_quantized, binary_quantized)
• muvera_fde
• clustered variants (clustered_8x, clustered_16x, clustered_32x, clustered_64x)

from common.qdrant_collections import recreate_colpali_muvera_collection

recreate_colpali_muvera_collection(
    qdrant,
    collection_name="colpali_perf_benchmark",
    colpali_dim=128,
    muvera_dim=MUVERA_DIM,
    enable_optimized=True,
    clustered_k_values=[8, 16, 32, 64],
)

2) Optimization Ideas (Intuition First)

All methods in this post follow the same two-stage idea:

1. Fast prefetch on compressed vectors to get candidate pages.
2. Full rerank with colpali_original to recover quality.

This is important because we are not throwing away ColPali scoring; we are using cheaper vectors as a candidate filter.

2.1 Dual-axis pooling

A ColPali page can be viewed as a patch grid (roughly 32 x 32 in this setup). Dual-axis pooling compresses this grid in two ways:

• row pooling: average along columns
• column pooling: average along rows

Intuition: rows and columns may preserve different layout cues (horizontal vs vertical structure), and combining both candidate lists should improve coverage.

2.2 Token clustering

Instead of averaging by spatial axis, token clustering groups similar patch tokens and keeps only centroid vectors.

def cluster_tokens_kmeans(image_tokens, k=16, n_init=5, random_state=42):
    from sklearn.cluster import KMeans
    k_eff = max(1, min(int(k), image_tokens.shape[0]))
    km = KMeans(n_clusters=k_eff, n_init=n_init, random_state=random_state)
    km.fit(image_tokens)
    return km.cluster_centers_.astype("float32")

Intuition: semantic clusters can preserve diverse content (charts, text blocks, logos, figures) better than simple averaging.

Two-stage retrieval flow

3) Experiments and Results

3.1 First experiment: dual-axis pooling + rerank

Method:

• prefetch from row_pooled and column_pooled
• union candidates
• rerank with colpali_original

The outcome:

Table : Dual-axis pooling + rerank results

This was the first major surprise: dual-axis was not better than row-only in practice.

Candidate overlap analysis explained why:

• row vs column prefetch candidate sets had mean Jaccard overlap = 0.9958
• they were almost the same list most of the time

So the “two views should add diversity” hypothesis mostly failed on this dataset.

3.2 Token clustering approach

Method:

• build clustered page vectors (clustered_8x, clustered_16x, clustered_32x, clustered_64x)
• prefetch on clustered vectors
• rerank with full colpali_original

Results:

• 32× compression offers the optimal balance — near-baseline quality with 4× speedup
• 16× compression sacrifices more accuracy (Recall@10: 0.732) with similar latency
• 8× compression shows significant quality drop (Recall@10: 0.690), suggesting over-compression

The tradeoff was much better than pooling: quality stayed close while latency dropped significantly.

3.3 Ablation studies

I then ran controlled sweeps on prefetch depth, cluster K, axis choices, and reranker.

Table : Ablation study

Takeaway: deeper prefetch helps dual-axis a bit, but it never catches up. Clustered variants remain strong and fast even at low/moderate prefetch.

3.4 Final comparison

clustered_64x_rerank was the most balanced winner:

• slightly better recall/NDCG (0.631) than baseline(0.630),
• roughly 2x lower latency (23.7 ms) than full baseline (49.7 ms)in this run.

Quality vs latency scatter plot

4) Discussion: What Worked (and What Didn’t)

Why dual-axis pooling struggled

Pooling is aggressive averaging. It can remove token-level differences that MaxSim depends on, especially for visually nuanced document patches.

The overlap result (~0.996 Jaccard) is the key clue: row and column branches did not produce complementary candidates. So dual-axis was mostly “same candidates twice,” plus extra overhead.

Why clustering worked

Clustering keeps representative token centroids instead of flattening everything by axis. That preserved more semantic diversity while reducing vector count sharply.

A practical way to read it:

• original: ~1,030 vectors/page
• clustered_64x: 64 vectors/page
• compression ratio: ~16x
• quality: near-baseline (or slightly better in this setup)
• Storage and memory implications

This project did not include absolute RAM/disk deltas per method, but vector-count compression gives a strong proxy:

• fewer vectors per page means smaller multivector payloads,
• smaller payloads reduce index/storage pressure and improve query-time efficiency.