ColPali Optimization Project

This repository contains the code used to study and improve ColPali retrieval in Qdrant. The main goal is to reduce latency while keeping retrieval quality close to the full baseline.

ColPali is a late-interaction model that turns each document image into many patch token vectors. Search then compares query tokens with image tokens using MaxSim scoring.

Project Structure

• common/: shared Python utilities for collection setup, retrieval methods, evaluation metrics, and plotting.
• data/: main evaluation dataset assets (images/, queries.json, relevance.csv) and the download script.
• data_preso/: smaller presentation dataset with the same layout as data/.
• local_model_demo/: local embedding flow (GPU model execution with ColQwen3).
• server_engine_demo/: server-based embedding flow (calls a running ColPali server).
• optimization_experiments/: core experiments for pooling, clustering, ablations, and final comparison.

Setup

1. Install dependencies:

pip install -r requirements.txt

2. Create a .env file in the project root:

COLPALI_SERVER_URL=http://localhost:8000

3. Start Qdrant (used by notebooks).
   Default URLs in this project:

• http://localhost:6333 (demo notebooks)
• http://localhost:6335 (optimization notebooks)

4. Download and prepare the dataset:

python data/download_vidore.py --max-pages 500

Data Setup

The data comes from vidore/docvqa_test_subsampled on Hugging Face.

Expected output after running the download script:

• data/images/*.png: document page images
• data/queries.json: query list
• data/relevance.csv: query-to-image relevance labels

Typical dataset size used in experiments:

• 500 unique page images
• 500 query rows
• 1 relevance label per query in this setup

queries.json format

[
  {"id": "q1", "text": "coffee mug on a desk"},
  {"id": "q2", "text": "transformer architecture diagram"}
]

relevance.csv format

query_id,image_id,relevance
q1,image_0001,1
q1,image_0023,1
q2,image_0104,2

Note: image_id is expected to match the image filename (without extension), unless a custom resolver is used.

Optimization Experiments

All main experiment notebooks are in optimization_experiments/.

00_setup_collection.ipynb

Purpose:

• Creates the Qdrant collection with baseline and optimized vector slots.
• Optionally re-indexes image embeddings into all vector fields.

Named vectors created:

• colpali_original
• muvera_fde
• scalar_quantized
• binary_quantized
• hierarchical_2x
• hierarchical_4x
• row_pooled
• column_pooled
• clustered_8x, clustered_16x, clustered_32x, clustered_64x

Observed run from notebook output:

• 500 images indexed in about 2182.8 seconds (~36.4 minutes)

01_dual_axis_pooling.ipynb

Purpose:

• Tests pooled prefetch strategies with full ColPali rerank.
• Compares row_pooled, column_pooled, and dual-axis prefetch.

Key result:

• Full baseline recall is 0.764
• Dual-axis rerank recall is 0.422
• Row-only rerank recall is 0.420
• Column-only rerank recall is 0.384

Candidate overlap check:

• Row vs column prefetch candidate lists have mean Jaccard overlap 0.9958
• This means both lists are almost the same for most queries

02_token_clustering.ipynb

Purpose:

• Clusters patch tokens with KMeans and uses cluster centroids for prefetch.
• Reranks final candidates with full colpali_original.

Cluster settings:

• clustered_8x, clustered_16x, clustered_32x

Key result from notebook:

• Baseline: recall 0.764, latency 194.76 ms
• clustered_32x_rerank: recall 0.754, latency 48.41 ms

This is close recall with much lower latency.

03_ablation_studies.ipynb

Purpose:

• Runs controlled sweeps on important retrieval settings.

Ablations covered:

• Prefetch depth sweep (20, 30, 50, 100)
• Single-axis vs dual-axis pooling
• Cluster K sweep (8, 16, 32, 64)
• Reranker choice (colpali_original, binary_quantized, hierarchical_2x)

Important findings:

• Dual-axis pooling stays low recall even with deeper prefetch (up to 0.452 at prefetch 100)
• Clustered methods keep high recall and low latency
• clustered_64x reaches recall around or above baseline in this setup

04_final_comparison.ipynb

Purpose:

• Compares target approaches side by side on shared metrics.
• Produces a final recommendation.

Final table from notebook output:

Notebook recommendation:

• clustered_64x_rerank
• Reason: best NDCG among eligible methods, recall matches baseline, and latency is about 2x faster than full baseline.

Shared Utility Module

optimization_experiments/pooling_utils.py contains the helper functions used in multiple notebooks:

• safe_row_column_pool: computes row/column pooled vectors when patch grid metadata is valid.
• hierarchical_pool: groups adjacent tokens by fixed size (for example 2x or 4x).
• cluster_tokens_kmeans: runs KMeans on image tokens and returns centroid tokens.
• candidate_ids: extracts candidate IDs from ranked results.
• jaccard: computes candidate-set overlap.

Other Notes

• Main Python dependencies are listed in requirements.txt.
• .env is ignored by git (.gitignore) and should stay local.
• Image files (*.png) are also ignored, so downloaded data is not committed by default.