Vision-Based Document Retrieval

Why spend months building a document parsing pipeline when a vision model can just look at the page?

Vec, Sharp-Eyed AI Agent

1. The Problem with Text-Only Document Retrieval

The standard RAG pipeline for documents follows a predictable sequence: parse the document (PDF, Word, HTML) into text, chunk the text, embed the chunks, and store them in a vector database. This pipeline works well for text-heavy documents but struggles with visually rich content. The failure modes are systematic and well-documented:

• Tables lose structure: A PDF table becomes a jumbled sequence of cell values after text extraction. The spatial relationships between headers and data cells, which carry the meaning, are destroyed.
• Charts become invisible: Bar charts, line graphs, and diagrams contain information that never appears in extracted text. A pie chart showing market share is simply absent from the text index.
• Layout carries meaning: In financial reports, the position of numbers on the page (which column, which row) determines what they mean. Flat text extraction discards this positional information.
• OCR introduces errors: Scanned documents require OCR, which introduces character-level errors that accumulate into retrieval failures. A misspelled company name or garbled number will not match any query.
• Multi-language documents: Documents mixing scripts (English headers with Japanese body text, Arabic annotations on English diagrams) are particularly fragile under text extraction.

These failures are not edge cases. In enterprise settings, a significant fraction of knowledge lives in slide decks, scanned forms, technical drawings, and formatted reports where layout and visual elements carry essential information. The traditional text pipeline silently drops this information, and no amount of embedding model improvement can recover what was never extracted.

2. ColPali: The Architecture

ColPali (2024) is the model that brought vision-based document retrieval from research curiosity to practical tool. Its name combines "Col" (from ColBERT's late interaction mechanism, introduced in Section 19.1) with "Pali" (from Google's PaLI vision-language model). The architecture has two key components: a vision-language backbone that produces per-patch embeddings from document images, and a late interaction scoring mechanism that computes fine-grained relevance between query tokens and image patches.

How ColPali Processes a Page

Given a document page rendered as an image (typically 448x448 or higher resolution), ColPali processes it through a vision transformer that divides the image into a grid of patches (for example, 32x32 patches). Each patch produces a 128-dimensional embedding vector. A single page therefore produces approximately 1,024 embedding vectors, each representing a small region of the page. These patch embeddings encode both the textual content visible in that region and the visual context (layout, formatting, surrounding elements).

The query side is simpler: a text query is tokenized and each token is projected into the same 128-dimensional space. The relevance score between a query and a page is computed using MaxSim (Maximum Similarity), the same late interaction mechanism from ColBERT.

MaxSim: Late Interaction Scoring

The MaxSim scoring function is the key to ColPali's retrieval quality. For a query q with tokens q₁, $q_{2}$, ..., $q_{n}$ and a document page with patch embeddings p₁, $p_{2}$, ..., $p_{m}$, the relevance score is:

In words: for each query token, find the document patch that is most similar to it, then sum these maximum similarities across all query tokens. This allows different parts of the query to match different regions of the page. The token "revenue" might match a cell in a table, while "chart" matches a visual element in the corner, and "Q3" matches a heading at the top. Single-vector retrieval (like standard dense retrieval) compresses the entire page into one vector and cannot capture this multi-region matching.

A small numeric example with 2 query tokens and 3 document patches makes this concrete:

# MaxSim scoring: numeric walkthrough
import numpy as np

# Similarity matrix: sim(q_i, p_j) for 2 query tokens, 3 patches
sim = np.array([
 [0.8, 0.3, 0.1], # query token "revenue" similarities to patches
 [0.2, 0.5, 0.9], # query token "chart" similarities to patches
])

max_per_query = sim.max(axis=1) # [0.8, 0.9]
score = max_per_query.sum() # 0.8 + 0.9 = 1.7
print(f"MaxSim per query token: {max_per_query}") # [0.8, 0.9]
print(f"Total MaxSim score: {score}") # 1.7
# "revenue" matched patch 0 best; "chart" matched patch 2 best.

3. ColQwen and the Evolving Model Family

ColPali was the first model to demonstrate the approach, but the architecture generalizes to any vision-language backbone. ColQwen2 (2024) replaces PaLI with the Qwen2-VL vision-language model, achieving stronger performance on the ViDoRe benchmark. The key differences: Code Fragment 19.5.2 below puts this into practice.

• Higher resolution: ColQwen2 processes images at higher resolution with dynamic resolution support, producing more patch embeddings per page and better capturing fine-grained text and small visual elements.
• Stronger multilingual support: Qwen2-VL's multilingual pretraining transfers to document retrieval, making ColQwen2 significantly better on non-English documents.
• Improved training data: ColQwen2 was trained on a larger and more diverse set of query-page pairs, including financial documents, scientific papers, slide decks, and forms.

# implement embed_pages, embed_queries, maxsim_score
# Key operations: training loop, embedding lookup, results display
import torch
from colpali_engine.models import ColQwen2, ColQwen2Processor
from PIL import Image
from typing import List

# Load model and processor
model_name = "vidore/colqwen2-v1.0"
model = ColQwen2.from_pretrained(
 model_name,
 torch_dtype=torch.float16,
 device_map="auto",
)
processor = ColQwen2Processor.from_pretrained(model_name)

def embed_pages(images: List[Image.Image]) -> torch.Tensor:
 """
 Embed document page images into multi-vector representations.

 Returns tensor of shape (n_pages, n_patches, embedding_dim).
 """
 batch = processor.process_images(images).to(model.device)
 with torch.no_grad():
 embeddings = model(**batch)
 return embeddings # (n_pages, n_patches, 128)

def embed_queries(queries: List[str]) -> torch.Tensor:
 """
 Embed text queries into multi-vector representations.

 Returns tensor of shape (n_queries, n_tokens, embedding_dim).
 """
 batch = processor.process_queries(queries).to(model.device)
 with torch.no_grad():
 embeddings = model(**batch)
 return embeddings # (n_queries, n_tokens, 128)

def maxsim_score(
 query_embeddings: torch.Tensor,
 page_embeddings: torch.Tensor,
) -> torch.Tensor:
 """
 Compute MaxSim scores between queries and pages.

 Args:
 query_embeddings: (n_queries, n_tokens, dim)
 page_embeddings: (n_pages, n_patches, dim)

 Returns:
 Scores tensor of shape (n_queries, n_pages)
 """
 scores = torch.einsum(
 "qtd,npd->qntp", query_embeddings, page_embeddings
 )
 # For each query token, take max over all patches
 max_over_patches = scores.max(dim=-1).values # (q, n, t)
 # Sum over query tokens
 return max_over_patches.sum(dim=-1) # (q, n)

# Example usage
pages = [
 Image.open("financial_report_p1.png"),
 Image.open("financial_report_p2.png"),
 Image.open("org_chart.png"),
 Image.open("product_roadmap.png"),
]
queries = ["Q3 revenue breakdown by region", "engineering team structure"]

page_embs = embed_pages(pages)
query_embs = embed_queries(queries)
scores = maxsim_score(query_embs, page_embs)

# Rank pages for each query
for i, query in enumerate(queries):
 ranked = scores[i].argsort(descending=True)
 print(f"\nQuery: '{query}'")
 for rank, idx in enumerate(ranked[:3]):
 print(f" Rank {rank+1}: Page {idx.item()+1} (score: {scores[i][idx]:.2f})")

4. The ViDoRe Benchmark

The Visual Document Retrieval Benchmark (ViDoRe) was introduced alongside ColPali to provide a standardized evaluation for vision-based document retrieval. It covers a diverse set of document types and languages, making it the primary benchmark for comparing approaches. The benchmark includes:

On ViDoRe, ColQwen2 achieves NDCG@5 scores of 85 to 92% across datasets, compared to 60 to 75% for traditional text-extraction pipelines using the same queries. The largest gains appear on documents with complex visual elements: infographics (+25 points), financial tables (+18 points), and charts (+22 points). On purely text-heavy documents with clean layouts, the advantage narrows to 3 to 5 points, confirming that vision-based retrieval is most valuable when documents are visually rich.

5. Two-Stage Retrieval Pipeline

In production, pure ColPali/ColQwen retrieval can be expensive because late interaction requires storing and scoring against all patch embeddings. A practical architecture uses two stages: a fast first stage for candidate generation, followed by ColPali rescoring on the top candidates. Code Fragment 19.5.3 below demonstrates this approach.

# Define PageResult, TwoStageRetriever; implement __init__, index_page, retrieve
# Key operations: embedding lookup, results display, retrieval pipeline
import numpy as np
from typing import List, Dict, Tuple
from dataclasses import dataclass

@dataclass
class PageResult:
 page_id: str
 score: float
 source_doc: str
 page_number: int

class TwoStageRetriever:
 """
 Two-stage retrieval: fast first stage + ColQwen2 rescoring.

 Stage 1: Use single-vector embeddings (mean-pooled patch embeddings)
 stored in a standard vector DB for fast ANN search.
 Stage 2: Rerank top-K candidates using full MaxSim scoring
 against stored patch embeddings.
 """

 def __init__(self, model, processor, vector_db, patch_store):
 self.model = model
 self.processor = processor
 self.vector_db = vector_db # Standard vector DB (single vectors)
 self.patch_store = patch_store # Storage for full patch embeddings

 def index_page(self, page_id: str, image, metadata: dict):
 """Index a page for both retrieval stages."""
 # Generate patch embeddings
 batch = self.processor.process_images([image]).to(self.model.device)
 with torch.no_grad():
 patch_embs = self.model(**batch)[0] # (n_patches, dim)

 # Stage 1: Store mean-pooled vector in standard vector DB
 mean_vector = patch_embs.mean(dim=0).cpu().numpy()
 self.vector_db.upsert(
 ids=[page_id],
 vectors=[mean_vector.tolist()],
 metadata=[metadata],
 )

 # Stage 2: Store full patch embeddings for rescoring
 self.patch_store.save(page_id, patch_embs.cpu())

 def retrieve(
 self,
 query: str,
 first_stage_k: int = 100,
 final_k: int = 10,
 ) -> List[PageResult]:
 """
 Two-stage retrieval with ColQwen2 rescoring.

 Args:
 query: Text query
 first_stage_k: Candidates from first stage
 final_k: Final results after rescoring
 """
 # Stage 1: Fast candidate generation
 query_batch = self.processor.process_queries([query])
 query_batch = query_batch.to(self.model.device)
 with torch.no_grad():
 query_embs = self.model(**query_batch)[0] # (n_tokens, dim)

 # Mean-pool query for first-stage ANN search
 query_vector = query_embs.mean(dim=0).cpu().numpy()
 candidates = self.vector_db.query(
 vector=query_vector.tolist(),
 top_k=first_stage_k,
 )

 # Stage 2: MaxSim rescoring on candidates
 results = []
 for candidate in candidates:
 page_patches = self.patch_store.load(candidate.id)
 page_patches = page_patches.to(self.model.device)

 # MaxSim: for each query token, max similarity over patches
 sim_matrix = torch.matmul(
 query_embs, page_patches.T
 ) # (n_tokens, n_patches)
 max_sims = sim_matrix.max(dim=1).values # (n_tokens,)
 score = max_sims.sum().item()

 results.append(PageResult(
 page_id=candidate.id,
 score=score,
 source_doc=candidate.metadata.get("source", ""),
 page_number=candidate.metadata.get("page", 0),
 ))

 # Sort by MaxSim score and return top results
 results.sort(key=lambda r: r.score, reverse=True)
 return results[:final_k]

# Usage pattern
# retriever = TwoStageRetriever(model, processor, vector_db, patch_store)
# results = retriever.retrieve("quarterly revenue by product line", final_k=5)
# for r in results:
# print(f" {r.source_doc} p.{r.page_number}: {r.score:.1f}")

6. When to Use Vision-Based Retrieval

Vision-based retrieval is not a universal replacement for text-based retrieval. It excels in specific scenarios and introduces tradeoffs that practitioners should evaluate carefully.

Strong Use Cases

• Slide decks and presentations: Information is inherently visual (diagrams, formatted text boxes, images). Text extraction loses most of the content.
• Financial documents: Tables, charts, and formatted reports where layout encodes meaning.
• Scanned documents: Handwritten notes, historical records, forms with checkboxes. Vision models bypass OCR entirely.
• Technical drawings: Engineering diagrams, floor plans, circuit schematics where the visual content is primary.
• Multilingual document collections: Vision models handle mixed scripts naturally without language-specific text extraction.

Cases Where Text Retrieval Is Better

• Long text documents: A 50-page legal contract is better served by text chunking than by scoring 50 page images independently.
• Keyword-heavy queries: Exact term matching (case numbers, product codes, specific names) is faster and more reliable with text-based BM25.
• Very large corpora: The storage and computation costs of late interaction scale linearly with corpus size. At tens of millions of pages, the cost may be prohibitive without aggressive quantization.

7. Integration with RAG Pipelines

Vision-based retrieval slots into a RAG pipeline at the retrieval stage, but the downstream steps change. After retrieving relevant pages as images, you pass the page images directly to a multimodal LLM (GPT-4o, Claude, Gemini) rather than passing extracted text. The LLM receives the original page image and the question, preserving all visual context that would be lost in text extraction.

This creates a clean pipeline: render pages as images, embed with ColQwen2, store in a vector database, retrieve top pages for a query, send page images to a multimodal LLM for answer generation. The entire document parsing, chunking, and metadata extraction pipeline is replaced by a single rendering step (PDF to image), which is fast, deterministic, and lossless.