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

arXiv: https://arxiv.org/abs/2609.07262 · HTML: https://mlsanity.com/paper/2609.07262/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 39 references that are in the map.
- First posted on arXiv 7 Sep 2026.

## Builds on (most cited first, 27 of 27)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Sigmoid Loss for Language Image Pre-Training](https://mlsanity.com/paper/2303.15343/) (2023)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [DocVQA: A Dataset for VQA on Document Images](https://mlsanity.com/paper/2007.00398/) (2020)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Passage Re-ranking with BERT](https://mlsanity.com/paper/1901.04085/) (2019)
- [LayoutLM: Pre-training of Text and Layout for Document Image Understanding](https://mlsanity.com/paper/1912.13318/) (2019)
- [InfographicVQA](https://mlsanity.com/paper/2104.12756/) (2021)
- [OCR-free Document Understanding Transformer](https://mlsanity.com/paper/2111.15664/) (2021)
- [SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking](https://mlsanity.com/paper/2107.05720/) (2021)
- [VisRAG: Vision-based Retrieval-augmented Generation on Multi-modality Documents](https://mlsanity.com/paper/2410.10594/) (2024)
- [ColPali: Efficient Document Retrieval with Vision Language Models](https://mlsanity.com/paper/2407.01449/) (2024)
- [Unifying Multimodal Retrieval via Document Screenshot Embedding](https://mlsanity.com/paper/2406.11251/) (2024)
- [PLAID: An Efficient Engine for Late Interaction Retrieval](https://mlsanity.com/paper/2205.09707/) (2022)
- [Towards Complex Document Understanding By Discrete Reasoning](https://mlsanity.com/paper/2207.11871/) (2022)
- [Rethinking the Role of Token Retrieval in Multi-Vector Retrieval](https://mlsanity.com/paper/2304.01982/) (2023)
- [Reducing the Footprint of Multi-Vector Retrieval with Minimal Performance Impact via Token Pooling](https://mlsanity.com/paper/2409.14683/) (2024)
- [Introducing Neural Bag of Whole-Words with ColBERTer: Contextualized Late Interactions using Enhanced Reduction](https://mlsanity.com/paper/2203.13088/) (2022)
- [QuoTA: Query-oriented Token Assignment via CoT Query Decouple for Long Video Comprehension](https://mlsanity.com/paper/2503.08689/) (2025)
- [ModernVBERT: Towards Smaller Visual Document Retrievers](https://mlsanity.com/paper/2510.01149/) (2025)
- [A Study on Token Pruning for ColBERT](https://mlsanity.com/paper/2112.06540/) (2021)
- [PixelPrune: Pixel-Level Adaptive Visual Token Reduction via Predictive Coding](https://mlsanity.com/paper/2604.00886/) (2026)

## Abstract (arXiv)

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.