Visual document search works better with smart token selection

Query-Aware Token Budgeting for Efficient Late-Interaction Visual Document Retrieval

Information RetrievalArtificial Intelligence

Summary

Searching through images of documents usually involves breaking pages into many small parts called tokens, which helps find details but requires a lot of storage and time. The authors explored a way to first pick a small set of candidate pages quickly, and then decide which tokens to focus on based on the actual search question. They found that choosing tokens with this query-aware method recovers much more relevant information than just picking tokens without knowing the query. This approach improves the quality of search results while keeping the process efficient. It shows that adapting the search to the question is better than using a fixed way to summarize documents.

visual document retrievaltoken embeddingsquery-aware selectionlate-interaction modelsMaxSim coveragesubmodularitycandidate generationrerankingnormalized discounted cumulative gain (NDCG)pooling

Authors

PS Rishi, Rajeev Ranjan Dwivedi, Vinod K Kurmi

Abstract

Late-interaction visual document retrievers preserve fine-grained page evidence by storing many token embeddings per page, but the resulting storage and query-time interaction costs make large-scale deployment expensive. Pooling document tokens before indexing offers a natural remedy, yet static pooling must decide which visual evidence to preserve before the query is known. We study an alternative: a heavily compressed hot-path index generates candidates, after which query-aware token budgeting operates on the original token sets of the shortlisted pages. We formulate this stage-two selection as a budgeted MaxSim coverage problem, show that a clipped version is monotone submodular, and compare coverage-only, cluster-guided, token-wise, and marginal-gain policies. On ten ViDoRe tasks with ColModernVBERT, direct static pooling reduces macro normalized discounted cumulative gain at rank five from 0.6309 without compression to 0.4738 at a thirty-two-fold pool factor. Under the same candidate-generation regime and a pool-factor-eight-equivalent reranking budget, token top-k recovers 93.93 percent of the full-token score, while greedy marginal-gain selection recovers 98.39 percent. Held-out and leave-one-dataset-out evaluations yield positive greedy improvements over token top-k on every dataset. The latency analysis reveals two useful operating points: token top-k for interactive retrieval and the naive greedy implementation as a quality upper envelope. Together, these results show that late-interaction visual retrieval benefits from query-aware allocation rather than query-agnostic pooling alone.