Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Col-Bandit: Zero-Shot Query-Time Pruning for Late-Interaction Retrieval

Roi Pony, Adi Raz, Oshri Naparstek, Idan Friedman, Udi Barzelay

TL;DR

Col-Bandit tackles the high query-time cost of ColBERT-style late-interaction retrieval by adaptively pruning MaxSim computations at run time. By reframing reranking as a finite-population Top-$K$ identification problem and applying a LUCB-style refinement with variance-adaptive Serfling bounds, it reveals only a subset of token-level interactions until the Top-$K$ set is certified with a tunable relaxation parameter $\alpha_{\mathrm{ef}}$. The method is a drop-in layer that requires no index changes or offline retraining, and empirical results on BEIR and REAL-MM-RAG show up to about a 5x reduction in MaxSim FLOPs while preserving ranking fidelity and end-task performance. These findings suggest that dense late-interaction scoring contains substantial redundancy that can be efficiently pruned at query time, enabling practical, scalable deployment of high-accuracy multi-vector retrievers. The work provides a principled, theory-backed, and configurable framework for adaptive, on-the-fly pruning in dense, token-level retrieval.

Abstract

Multi-vector late-interaction retrievers such as ColBERT achieve state-of-the-art retrieval quality, but their query-time cost is dominated by exhaustively computing token-level MaxSim interactions for every candidate document. While approximating late interaction with single-vector representations reduces cost, it often incurs substantial accuracy loss. We introduce Col-Bandit, a query-time pruning algorithm that reduces this computational burden by casting reranking as a finite-population Top-$K$ identification problem. Col-Bandit maintains uncertainty-aware bounds over partially observed document scores and adaptively reveals only the (document, query token) MaxSim entries needed to determine the top results under statistical decision bounds with a tunable relaxation. Unlike coarse-grained approaches that prune entire documents or tokens offline, Col-Bandit sparsifies the interaction matrix on the fly. It operates as a zero-shot, drop-in layer over standard multi-vector systems, requiring no index modifications, offline preprocessing, or model retraining. Experiments on textual (BEIR) and multimodal (REAL-MM-RAG) benchmarks show that Col-Bandit preserves ranking fidelity while reducing MaxSim FLOPs by up to 5$\times$, indicating that dense late-interaction scoring contains substantial redundancy that can be identified and pruned efficiently at query time.

Col-Bandit: Zero-Shot Query-Time Pruning for Late-Interaction Retrieval

TL;DR

Col-Bandit tackles the high query-time cost of ColBERT-style late-interaction retrieval by adaptively pruning MaxSim computations at run time. By reframing reranking as a finite-population Top- identification problem and applying a LUCB-style refinement with variance-adaptive Serfling bounds, it reveals only a subset of token-level interactions until the Top- set is certified with a tunable relaxation parameter . The method is a drop-in layer that requires no index changes or offline retraining, and empirical results on BEIR and REAL-MM-RAG show up to about a 5x reduction in MaxSim FLOPs while preserving ranking fidelity and end-task performance. These findings suggest that dense late-interaction scoring contains substantial redundancy that can be efficiently pruned at query time, enabling practical, scalable deployment of high-accuracy multi-vector retrievers. The work provides a principled, theory-backed, and configurable framework for adaptive, on-the-fly pruning in dense, token-level retrieval.

Abstract

Multi-vector late-interaction retrievers such as ColBERT achieve state-of-the-art retrieval quality, but their query-time cost is dominated by exhaustively computing token-level MaxSim interactions for every candidate document. While approximating late interaction with single-vector representations reduces cost, it often incurs substantial accuracy loss. We introduce Col-Bandit, a query-time pruning algorithm that reduces this computational burden by casting reranking as a finite-population Top- identification problem. Col-Bandit maintains uncertainty-aware bounds over partially observed document scores and adaptively reveals only the (document, query token) MaxSim entries needed to determine the top results under statistical decision bounds with a tunable relaxation. Unlike coarse-grained approaches that prune entire documents or tokens offline, Col-Bandit sparsifies the interaction matrix on the fly. It operates as a zero-shot, drop-in layer over standard multi-vector systems, requiring no index modifications, offline preprocessing, or model retraining. Experiments on textual (BEIR) and multimodal (REAL-MM-RAG) benchmarks show that Col-Bandit preserves ranking fidelity while reducing MaxSim FLOPs by up to 5, indicating that dense late-interaction scoring contains substantial redundancy that can be identified and pruned efficiently at query time.
Paper Structure (47 sections, 25 equations, 8 figures, 9 tables, 3 algorithms)

This paper contains 47 sections, 25 equations, 8 figures, 9 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Background and Related Work
  4. Preliminaries: Late Interaction Retrieval
  5. ColBERT Late Interaction Scoring.
  6. Top-$K$ Retrieval.
  7. Index-Time vs. Query-Time.
  8. Related Work
  9. Index-Time Compression & Token Pruning.
  10. Efficient Systems & Bound-Based Pruning.
  11. MaxSim-Level Pruning (Our Approach).
  12. Finite-Population Bandits and Top-$k$ Arm Identification.
  13. Problem Formulation
  14. The MaxSim Matrix and Observation Model
  15. Mapping to Finite-Population Bandits
  16. ...and 32 more sections

Figures (8)

  • Figure 1: Schematic of Col-Bandit. Given a query (e.g., "human mobility...") and a set of candidate documents (e.g., Nature, Auto), the goal is to identify the Top-$2$ relevant documents. (A) Full ColBERT determines the exact score for every document by summing all interaction cells (MaxSims), requiring 100% of the compute budget. (B) Col-Bandit approximates these sums using partial cell observations. By adaptively revealing informative cells (green) and skipping others (hatched), it maintains confidence intervals for the total score. The algorithm terminates as soon as a positive separation gap emerges: the Lower Bound of the weakest winner (Sports) is strictly higher than the Upper Bound of the strongest loser (Auto). This enables the identification of the correct Top-$K$ ranking while saving 60% of the query-time computations.
  • Figure 2: Cost-Accuracy trade-off for Col-Bandit compared to Random Reveal (Doc-Uniform) and Greedy Top-Margin (Doc-TopMargin) across three retrieval settings (text and multimodal). Each star marker denotes a Col-Bandit operating point obtained by sweeping the relaxation parameter $\alpha_{\mathrm{ef}}$. The top-right corner (Overlap@5 = 1.0, Cost = 100%) corresponds to full exhaustive scoring.
  • Figure 3: Taxonomy of efficient late-interaction retrieval. Methods are classified by when they prune (index-time vs. query-time) and what they prune. Col-Bandit is the first to dynamically prune the atomic interaction matrix $H$ during query-time scoring.
  • Figure 4: Exploration Strategy Ablation. Trade-off on Jina ColBERTv2 / HotPotQA. The dynamic $\epsilon$-greedy policy (purple) consistently dominates static warm-up schedules (green), avoiding wasteful reveals on easy queries.
  • Figure 5: Effect of ANN-derived bounds. Col-Bandit (purple) outperforms the corresponding baseline (gray) in both settings: with retrieval bounds (solid) and without (dashed). Granite Vision Embedding / TechSlides.
  • ...and 3 more figures