Finding the Needle in a Haystack: Comparing Multi-Vector Search Strategies

In the development of advanced search and retrieval systems, moving from keyword matching to semantic understanding is a critical step. However, a key distinction exists between finding a relevant document and locating a specific piece of information within that document with precision. While there are techniques that perform well for retrieving documents, most of them work by extracting summarized semantic meaning of the document. This can be seen as these models trying to understand the “gist” of the documents. Both single-vector search and late-interaction approaches work well for these conditions with various tradeoffs involved. But there’s another type of problem where the goal is not just to understand the overall topic of a document in general, but to also specifically account for the requested detail within the document. This “needle in a haystack” problem is a significant challenge, and addressing it is essential for building high-precision retrieval systems.

This guide provides a technical analysis of multi-vector search for high-precision information retrieval. We will examine various optimization strategies and analyze their performance. This guide should be seen as complementary to resources like the Answer.AI blog on ColBERT pooling , which explains how pooling strategies can be effective for document-level retrieval. Here, we will demonstrate why those same techniques can be counterproductive when precision at an intra-document, token level is the primary objective.

Find reproducible code here

The Dataset

This task is different from benchmarks like BEIR, which focus on text-based doc retrieval, finding the most relevant documents from a large collection. Here, we want intra-document localization, where the goal is to find a precise piece of information within a single, dense document, in a multimodal setting.

The Task: The Document Haystack Dataset

Our benchmark is built on the AmazonScience/document-haystack dataset, which contains 25 visually complex source documents (e.g., financial reports, academic papers). To create a rigorous test, our evaluation follows a per-document methodology:

• We process each of the 25 source documents independently.
• Table Creation: For a single source document (e.g., “AIG”), we ingest all pages from all of its page-length variants (from 5 to 200 pages long). This creates a temporary LanceDB table containing approximately 1,230 pages.
• The Task: We then query this table using a set of “needle” questions, where the goal is to retrieve the exact page number containing the answer. A successful retrieval means the correct page number is within the top K results.
• Target metric: We measure both retrieval accuracy (Hit@K) and the average search latency for each query against this table.

Dataset example The documents contain “text needles” like these

During evaluation, the queries processed are somewhat like this:

What is the secret currency in the document?
What is the secret object #3 in the document?

The intention of this task is to find the page which has the text needle that answers this questions

Models and Architectures

Our testbed includes a baseline single-vector model and a family of advanced multi-vector models.

Single-Vector (Bi-Encoder) Baseline: openai/clip-vit-base-patch32

A bi-encoder maps an entire piece of content (a query, a document page) to a single vector. The search process is simple: pre-compute one vector for every page, and at query time, find the page vector closest to the query vector.

• Strength: Speed and simplicity.
• Weakness: This creates an information bottleneck. All the nuanced details, keywords, and semantic relationships on a page must be compressed into a single, fixed-size vector. For finding a needle, this is like trying to describe a specific person’s face using only one word.

Multi-Vector (Late-Interaction) Models

Multi-vector models, pioneered by ColBERT, take a different approach. Instead of one vector per page, they generate a set of vectors for each page—one for every token (or image patch).

• Mechanism (MaxSim): The search process is more sophisticated. For each token in the query, the system finds the most similar token on the page. These maximum similarity scores are then summed up to get the final relevance score. This “late-interaction” preserves fine-grained, token-level details.
• The Models: We used several vision-language models adapted for this architecture, including ColPali, ColQwen2, and ColSmol. While their underlying transformer backbones differ, they all share the ColBERT philosophy of representing documents as a bag of contextualized token embeddings.

Different Retrieval Strategies Used

A full multi-vector search is powerful but computationally intensive. Here are five strategies for managing it, complete with LanceDB implementation details.

1. base: The Gold Standard (Full Multi-Vector Search)

This is the pure, baseline late-interaction search. It offers the highest potential for accuracy by considering every token.

LanceDB also integrates with ConteXtualized Token Retriever (XTR) , an advanced retrieval model that prioritizes the most semantically important document tokens during search. This integration enhances the quality of search results by focusing on the most relevant token matches. LanceDB Implementation:

import lancedb
import pyarrow as pa

# Schema for multi-vector data
# Assumes embeddings are 128-dimensional
schema = pa.schema([
    pa.field("page_num", pa.int32()),
    pa.field("vector", pa.list_(pa.list_(pa.float32(), 128)))
])

db = lancedb.connect("./lancedb")
tbl = db.create_table("document_pages_base", schema=schema)

# Ingesting multi-token embeddings for a page
# multi_token_embeddings is a NumPy array of shape (num_tokens, 128)
tbl.add([{"page_num": 1, "vector": multi_token_embeddings.tolist()}])

# Searching with a multi-token query
# query_multi_vector is also shape (num_query_tokens, 128)
results = tbl.search(query_multi_vector).limit(5).to_list()

2. flatten: Mean Pooling

This strategy “flattens” the set of token vectors into a single vector by averaging them. This transforms the search into a standard, fast approximate nearest neighbor (ANN) search.

LanceDB Implementation:

# Schema for single-vector data
schema_flat = pa.schema([
    pa.field("page_num", pa.int32()),
    pa.field("vector", pa.list_(pa.float32(), 128))
])
tbl_flat = db.create_table("document_pages_flat", schema=schema_flat)

# Ingesting the mean-pooled vector
mean_vector = multi_token_embeddings.mean(axis=0)
tbl_flat.add([{"page_num": 1, "vector": mean_vector.tolist()}])

# Searching with a single averaged query vector
query_mean_vector = query_multi_vector.mean(axis=0)
results = tbl_flat.search(query_mean_vector).limit(5).to_list()

3. max_pooling

This is a variation of flatten. max_pooling takes the element-wise max across all token vectors instead of the mean. The implementation is identical to flatten, just with a different aggregation method (.max(axis=0)).

4. flatten and multi-vector rerank: The Hybrid “Optimization”

This two-stage strategy aims for the best of both worlds. First, use a fast, pooled-vector search to find a set of promising candidates. Then, run the full, accurate multi-vector search on only those candidates.

LanceDB Implementation: This requires a table with two vector columns.

# Schema with both flat and multi-vector columns
schema_rerank = pa.schema([
    pa.field("page_num", pa.int32()),
    pa.field("vector_flat", pa.list_(pa.float32(), 128)),
    pa.field("vector_multi", pa.list_(pa.list_(pa.float32(), 128)))
])
tbl_rerank = db.create_table("document_pages_rerank", schema=schema_rerank)

# Ingest both vectors
tbl_rerank.add([
    {
        "page_num": 1,
        "vector_flat": multi_token_embeddings.mean(axis=0).tolist(),
        "vector_multi": multi_token_embeddings.tolist()
    }
])

# --- Two-Stage Search ---
# Stage 1: Fast search on the flat vector
query_flat = query_multi_vector.mean(axis=0)
candidates = tbl_rerank.search(query_flat, vector_column_name="vector_flat") \
                       .limit(100) \
                       .with_row_id(True) \
                       .to_pandas()

# Stage 2: Precise multi-vector search on candidates
candidate_ids = tuple(candidates["_rowid"].to_list())
final_results = tbl_rerank.search(query_multi_vector, vector_column_name="vector_multi") \
                          .where(f"_rowid IN {candidate_ids}") \
                          .limit(5) \
                          .to_list()

5. hierarchical token pooling: Compressing the Haystack

This is an indexing-time strategy that aims to reduce the storage footprint and computational cost of multi-vector search by reducing the number of vectors per document. Instead of using every token vector, it clusters semantically similar tokens together and replaces them with a single, averaged vector.

• Mechanism: For each document, it computes the similarity between all token vectors, performs hierarchical clustering to group them, and then mean-pools the vectors within each cluster. This results in a smaller, more compact set of token vectors representing the document.
• Goal: To reduce memory and disk usage while attempting to preserve the most important semantic information, potentially offering a middle ground between the high accuracy of base search and the speed of pooled methods.

LanceDB Implementation: The schema is identical to the base multi-vector search, but the data is pre-processed before ingestion.

from utils import pool_embeddings_hierarchical
import numpy as np

# Schema is the same as the base multi-vector schema
schema = pa.schema([
    pa.field("page_num", pa.int32()),
    pa.field("vector", pa.list_(pa.list_(pa.float32(), 128)))
])
tbl_hierarchical = db.create_table("document_pages_hierarchical", schema=schema)

# Pool the embeddings before ingestion
# multi_token_embeddings is a NumPy array of shape (num_tokens, 128)
pooled_embeddings = pool_embeddings_hierarchical(
    multi_token_embeddings,
    pool_factor=4  # Reduce vector count by a factor of 4
)

# Ingest the smaller set of multi-token embeddings
# pooled_embeddings is now shape (approx. num_tokens / 4, 128)
tbl_hierarchical.add([{"page_num": 1, "vector": pooled_embeddings.tolist()}])

# Search is identical to the base multi-vector search
results = tbl_hierarchical.search(query_multi_vector).limit(5).to_list()

The Results

For a “needle in a haystack” task, retrieval accuracy is the primary metric of success. The benchmark results reveal a significant performance gap between the full multi-vector search strategy and common optimization techniques.

Baseline Performance: Single-Vector Bi-Encoder

First, we establish a baseline using a standard single-vector bi-encoder model, openai/clip-vit-base-patch32. This represents a common approach to semantic search but, as the data shows, is ill-suited for this task’s precision requirements.

With a Hit@20 rate of just under 12%, the baseline model struggles to reliably locate the correct page. This performance level is insufficient for applications requiring high precision.

Multi-Vector Model Performance

We now examine the performance of multi-vector models using different strategies. The following table compares the base (full multi-vector), flatten (mean pooling), and rerank (hybrid) strategies across several late-interaction models.

The data shows a consistent pattern: the base strategy outperforms all other techniques. The flattned pooling and reranking strategies perform no better than the single-vector baseline. However, hierarchical token pooling seems like a decent alternative to base considering speed vs accuracy tradeoff. Let’s look at the numbers in detail.

In-Depth Analysis of Pooling Strategies

To further understand the failure of optimization techniques, we compared different methods for pooling token vectors into a single vector: mean (flatten), max.

All flattened pooling methods perform poorly, confirming that the aggregation of token vectors into a single representation loses the fine-grained detail required for this task.

Finding the Right Trade-Off

1. The Failure of Simple Pooling: The flatten (mean pooling) and max_pooling strategies fail to improve upon the baseline. This is because their aggressive compression destroys the essential localization signal. The resulting single vector represents the topic of the page, not the specific needle on it.

2. The Failure of flatten and multi-vector rerank: This hybrid strategy is the worst-performing of all. The reason is a fundamental flaw in its design for this task: the first stage uses a simple pooled vector to retrieve candidates. Since this pooling eliminates the localization signal, the initial candidate set is effectively random.

3. hierarchical token pooling: By clustering and pooling tokens at indexing time, it reduces the number of vectors per page (in our case, by a factor of 4). This intelligently compresses the data, while preserving enough token-level detail in multi-vector setting. It achieves a Hit@20 of 91.6%, only slightly behind the base strategy’s 95.5%, but is significantly faster.

4. base multi-vector Search: The vanilla, un-optimized base multi-vector search remains the most accurate strategy. Preserving every token vector provides the highest guarantee of finding the needle, but this comes at the highest computational cost.

Latency:

The “optimizations” are not all created equal. While simple pooling is fast, its inaccuracy makes it unusable. Hierarchical pooling, however, offers a compelling balance of speed and accuracy.

(Latency measured on NVIDIA H100 GPUs)

Practical Considerations:

The accuracy of base multi-vector search is impressive, but its computational intensity has historically limited its use. hierarchical token pooling as a viable strategy creates a new, practical sweet spot on the accuracy-latency curve, making high-precision search accessible for a wider range of applications.

Search Latency and Computational Complexity

As the benchmark data shows, the search latency for base multi-vector search is orders of magnitude higher than for single-vector (or pooled-vector) search. It’s important to note that the reported ~670ms latency is an average from per-document evaluations. In this benchmark, each of the 25 documents is processed independently. All pages from a single document’s variants (ranging from 5 to 200 pages) are ingested into a temporary table, resulting in a table size of approximately 1,230 rows (pages) per evaluation. The search is performed on this table, and then the table is discarded. This highlights a significant performance cost even on a relatively small, per-document scale. This stems from a fundamental difference in computational complexity:

• Modern ANN Search (for single vectors): Algorithms like HNSW (Hierarchical Navigable Small World) provide sub-linear search times, often close to O(log N), where N is the number of items in the index. This allows them to scale to billions of vectors with millisecond-level latency.
• Late-Interaction Search (Multi-Vector): The search process is far more intensive. For each query, it must compute similarity scores between query tokens and the tokens of many candidate documents. The complexity is closer to O(M * Q * D), where M is the number of candidate documents to score, Q is the number of query tokens, and D is the average number of tokens per document. Hierarchical token pooling directly attacks this problem by reducing D, leading to a significant reduction in search latency.

When to Use Multi-Vector Search

Given these constraints, the choice of strategy depends on the specific requirements of the application.

• For Maximum Precision (base): In domains where the cost of missing the needle is extremely high, the full base search is the most reliable option.
• For a Balance of Precision and Performance (hierarchical token pooling): This is the ideal choice for many applications. It makes high-precision search practical for larger datasets and more interactive use cases where the sub-second latency of the base search may be too high. It significantly lowers the barrier to entry for adopting multi-vector search. It should still not be seen as a drop-in replacement for ANN, as it still requires more computational resources than single-vector search.
• For General-Purpose Document Retrieval (flatten / single-vector): For large-scale retrieval where understanding the “gist” is sufficient or where in cases where large-context text-based models suffice, single-vector search remains the most practical and scalable solution.

Appendix: Full Benchmark Results

Update 1: Add token pooling method