PaperlayerView exact revision on arXiv

Source-linked AI summary

Col-Bandit: Query-Time Top-$K$ Estimation for Late-Interaction Retrieval

Roi Pony, Adi Raz Goldfarb, Oshri Naparstek, Idan Friedman, Udi Barzelay, Eli Schwartz

arXiv:2602.02827v2cs.IR

TL;DR

Exhaustive token-level MaxSim scoring makes late-interaction reranking expensive even though top-K identification rarely requires every matrix entry. Col-Bandit estimates the exhaustive top-K through query-time partial matrix observation and confidence-based elimination, achieving substantial compute and CPU savings while retaining high overlap; its formal guarantee and hardware scope are limited.

  • Problem

    Exhaustive late-interaction reranking computes every token interaction for every document, creating a query-time bottleneck despite many unnecessary cells.

  • Method

    Col-Bandit progressively reveals MaxSim cells, maintains finite-population Bernstein–Serfling bounds, and eliminates documents whose upper bounds cannot reach the top-K.

  • Results

    7.0×/4.6× server-CPU and 12.9×/10.3× Apple-M1-Max speedups preserve Overlap@5 ≥0.90 on every corpus tested while reducing MaxSim computation by up to ∼8×.

  • Takeaways & Limitations

    Col-Bandit is a drop-in query-time reranking layer that requires no retraining or index changes and exposes strong cost–fidelity trade-offs against non-adaptive baselines.

  • Takeaways & Limitations

    The δ-PAC guarantee applies at αef=1, whereas deployed αef=0.2 uses a simplified radius and int8 estimates without a formal certificate; wall-clock evaluation covers only two CPU classes.

Abstract

from arXiv · show

Multi-vector late-interaction retrievers such as ColBERT achieve state-of-the-art quality, but their query-time cost is dominated by exhaustively computing token-level MaxSim interactions for every candidate document. The MaxSim scores of $N$ candidates against $T$ query tokens form an $N\times T$ matrix whose row-sums are the late-interaction scores, and identifying the top-$K$ rarely requires every entry. We introduce Col-Bandit, a query-time estimator of the exhaustive-MaxSim top-$K$: it reveals matrix entries in batches, maintains a finite-population Bernstein-Serfling confidence interval on each candidate's score, and permanently drops any document whose upper bound falls below the $K$-th largest lower bound, computing only the cells needed to separate the top-$K$. A single relaxation knob $α_{\mathrm{ef}}\in(0,1]$ tunes the compute-fidelity trade-off. We deploy $α_{\mathrm{ef}}{=}0.2$, while $α_{\mathrm{ef}}{=}1$ admits a $δ$-PAC guarantee under a simplified radius. On BEIR and REAL-MM-RAG, Col-Bandit preserves $\geq 90\%$ fidelity to the exhaustive top-$5$ on every corpus while cutting MaxSim FLOPs by up to ${\sim}8\times$, for up to ${\sim}13\times$ single-thread CPU speedups across x86 and ARM. A drop-in reranking layer, it needs no retraining or index changes.

1 Introduction

Col-Bandit treats late-interaction reranking as progressive, query-time matrix completion, computing only the MaxSim cells needed to identify the exhaustive top-K. It uses confidence-based elimination and a tunable relaxation while preserving ranking fidelity without retraining or index changes.

  • Motivation: Exact late interaction evaluates token-level MaxSim operations for every document, making query-time computation a bottleneck.The method is motivated by avoiding work on documents already ruled out.
  • Approach: Col-Bandit estimates the exhaustive MaxSim top-K by progressively revealing matrix cells and using the query to decide which document tokens matter.This query-time signal is unavailable to index-time accelerators.
  • Approach: Col-Bandit uses empirical Bernstein–Serfling confidence bounds for finite-population Top-K identification.The bounds exploit document token sequences being finite.
  • Approach: αef=1 provides a δ-PAC certificate, while the deployed αef=0.2 is a calibrated relaxation for the cost–fidelity trade-off.The relaxation tightens the decision radius relative to the unshrunk form.
  • Results: 7.0×/4.6× server-CPU and 12.9×/10.3× Apple-M1-Max speedups preserve Overlap@5 ≥0.90 across tested corpora.The evaluation covers BEIR and REAL-MM-RAG, with no index modifications or retraining.

2 Background and Related Work

Late-interaction retrieval gains fine-grained matching from token embeddings but pays for exhaustive token-level scoring. Col-Bandit addresses this query-time cost by adaptively pruning MaxSim cells using statistical bounds, complementing index-time and other pruning approaches.

  • Late Interaction Retrieval: ColBERT represents queries and documents as token embeddings, capturing fine-grained semantic matches that single-vector representations can miss.This representation supports high-accuracy neural retrieval.
  • Late Interaction Retrieval: Late-interaction scores sum each query token’s maximum similarity to document tokens, with similarity typically computed by cosine.The token-level operations form the MaxSim scoring process.
  • Top-K Ranking: The ranking objective is to identify the K documents with the highest MaxSim scores.The search set may be the full corpus or an upstream candidate set.
  • Index-Time vs. Query-Time: Standard reranking evaluates all N × T MaxSim values, and this grid dominates query-time cost even after candidate retrieval.Index-time methods must commit before seeing the query, whereas query-time methods can use the actual query.
  • Related Work: Col-Bandit is presented as the first method to adaptively use the query to prune atomic MaxSim cells during query-time scoring.It is positioned as distinct from index-time accelerators that quantize, project, or prune representations before querying.
  • Related Work: Col-Bandit is orthogonal to index-time methods and can rerank candidate sets produced by them.It also brings statistical bound-based pruning to dense late interaction.
  • Top-K Identification: The method extends fixed-confidence Top-K identification to finite populations sampled without replacement, using Bernstein–Serfling concentration.Each MaxSim row is treated as a fixed set of deterministic values.

3 Problem Formulation

The paper formulates late-interaction reranking as finite-population Top-K identification over a partially observed MaxSim matrix. Col-Bandit reveals entries, maintains document bounds, and stops when those bounds separate winners from losers.

  • 3 Problem Formulation: Col-Bandit is formulated as a statistical estimator that outputs an estimated top-K set from the MaxSim matrix.Under αef=1, it recovers the true top-K with probability at least 1−δ.
  • 3 Problem Formulation: The formal problem is a fixed-confidence Multi-Armed Bandit problem over a sparsely observed, finite-population MaxSim matrix.The finite population consists of deterministic MaxSim values for each document’s query-token interactions.
  • 3.1 The MaxSim Matrix and Observation Model: The implicit MaxSim matrix H has N rows for documents and T columns for query tokens, with each entry representing a maximum token similarity.Each document’s total late-interaction score is the corresponding row-sum.
  • 3.1 The MaxSim Matrix and Observation Model: Full ColBERT scores every cell, whereas Col-Bandit reveals only a subset and skips the rest while tracking [LCBi, UCBi] for each document.The partial-observation view is the basis for reducing computation.
  • 3.1 The MaxSim Matrix and Observation Model: The top-K is identified once the weakest winner’s lower bound exceeds the strongest loser’s upper bound.This positive separation gap provides the stopping condition for partial observation.
  • 3.1 The MaxSim Matrix and Observation Model: At any step, observed cells define per-document revealed and unrevealed token sets, and revealing an unseen cell incurs unit cost.This makes the observation process explicit at the document-token level.
  • 3.1 The MaxSim Matrix and Observation Model: Coverage measures computational cost as the fraction of the N × T matrix that has been revealed.The algorithm’s cost is therefore tied to partial matrix exposure rather than only final reranking time.

4 Method: Col-Bandit

Col-Bandit estimates the exhaustive top-K by progressively revealing batched MaxSim cells, maintaining per-document bounds, and eliminating candidates that cannot remain competitive. It combines finite-population confidence intervals with batched rescoring and offers a certified αef=1 setting alongside a more aggressive relaxation.

  • Decision bounds: Col-Bandit maintains per-document lower and upper bounds, identifying the top-K when the weakest winner's lower bound exceeds the strongest loser's upper bound.The algorithm progressively drives this positive separation gap through successive reveal rounds.
  • Decision bounds: Global support bounds remain valid for bounded token-level similarities, while wider support preserves elimination safety for unnormalized embeddings at the cost of looser radii.The method uses simpler global support throughout rather than tighter per-cell bounds.
  • Decision bounds: αef=1 recovers the unshrunk empirical Bernstein–Serfling radius, whereas αef<1 tightens it and makes elimination more aggressive.The relaxation controls conservativeness of the decision interval.
  • Batched progressive elimination: Each round reveals B cells for every active document in one vectorized pass, refreshes bounds, and removes documents below the K-th largest lower bound.A safety margin M preserves up to K+M borderline survivors for exact rescoring on all T query tokens.
  • Certified and relaxed settings: αef=1 admits a δ-PAC guarantee, while the deployed αef=0.2 relaxation has no formal certificate but achieves at least 0.90 overlap with exhaustive top-K on every tested corpus.The guarantee follows from uniform-without-replacement reveals and simultaneous validity of the confidence intervals.

5 Experiments

Experiments evaluate Col-Bandit across BEIR and REAL-MM-RAG using coverage, top-K overlap, ranking quality, latency, scaling, and calibration analyses. The method achieves substantial compute and wall-clock savings while preserving near-exhaustive ranking fidelity across platforms, encoders, and complementary compression settings.

  • Experimental setup: Col-Bandit is evaluated on five BEIR corpora and four REAL-MM-RAG multimodal corpora using Full-MaxSim, Doc-Uniform, and Ball-carving baselines.Reported metrics include Overlap@K, coverage, wall-clock latency, and task-level ranking metrics.
  • Coverage savings: 14%/26%/22% coverage recovers 95% Overlap@5 for ColBERTv2/Jina/GVE, versus 98–100% for Doc-Uniform and 46–76% for Ball-carving.The 2–5× coverage gap holds across every corpus and at K=50.
  • Coverage savings: 9% coverage recovers 95% nDCG@5 on Jina and 13% on ColBERTv2, corresponding to 10.9× and 7.6× savings respectively.Recall@5 and MRR@5 show the same 13–26% coverage pattern for 95% retention on every corpus.
  • Wall-clock translation: 7.0× mean single-thread CPU-S speedup and 12.9×/10.3× CPU-M1 speedups are achieved at αef=0.2 while preserving Overlap@K ≥0.90 on every corpus.The CPU-S mean falls to 5.3× at K=100, while CPU-M1 reaches 9.0×/7.0× at K=100.
  • Sensitivity and scaling: 1.3× and 2.7× latency increases result from a 200× K jump on HotpotQA and SciDocs, respectively, demonstrating sub-linear scaling in K.The increase is 4.5× on the smaller ArguAna corpus because its active-set decisions amortize less effectively.
  • Composability: 26% coverage reaches 95% Overlap@5 on Jina, while Matryoshka dimension reduction preserves ≥0.95 overlap and leaves coverage essentially invariant.Col-Bandit also composes with Ward token pooling, with end-to-end loss empirically decomposed into pooling-induced and cell-skipping effects.
  • Calibration: αef=0.2 keeps 13–14% coverage at Overlap@5 ≥0.92, whereas αef=1 reaches Overlap@5=1.00 at 28–59% coverage.The knob therefore exposes a compute–fidelity trade-off, with αef=1 corresponding to the δ-PAC corner under the stated radius.

6 Conclusion

Col-Bandit accelerates query-time late-interaction reranking by selectively revealing MaxSim cells until the top-K set stabilizes. Across BEIR and REAL-MM-RAG, it reduces computation and latency while preserving high overlap with exhaustive reranking.

  • Col-Bandit selectively reveals MaxSim cells until the top-K set stabilizes.
  • Up to ∼8× lower MaxSim FLOPs yields 7.0×/4.6× server-CPU and 12.9×/10.3× Apple M1 Max wall-clock speedups.
  • Overlap@5 remains ≥0.90 on every tested corpus, reaching ≥0.96 on CPU-S except HotpotQA-500 K.
  • The αef knob controls the quality–compute trade-off, while αef=1 recovers a δ-PAC guarantee.
  • Col-Bandit is a drop-in reranking layer requiring no retraining or index changes.

Limitations

The main limitations concern hardware coverage, fidelity scope, the gap between deployment and formal theory, and unavailable Quora end-task metrics.

  • Wall-clock evaluation covers only an AVX2 server and a NEON laptop, without a GPU-native fused kernel.
  • Col-Bandit can recover the underlying exhaustive MaxSim top-K but cannot exceed the quality of that scorer.
  • The δ-PAC certificate applies only at αef=1; deployed αef=0.2 uses a simplified radius and int8-quantised estimates without formal certification.
  • Quora end-task metrics are unavailable because of a qrels/doc-id mapping issue, although Overlap@K remains unaffected.

Ethics Statement

This work accelerates an existing MaxSim retrieval primitive on public BEIR and REAL-MM-RAG benchmarks. It introduces no new datasets, models, or human-subjects data, and adds no new model capability.

  • The work accelerates the existing MaxSim retrieval primitive on public BEIR and REAL-MM-RAG benchmarks.
  • It introduces no new datasets, models, or human-subjects data.
  • Its only dual-use consideration is generic to retrieval-efficiency methods that speed search over an already indexed corpus.
  • No new model capability is introduced.

A.2 Datasets and Models

The evaluation spans text and multimodal retrieval datasets, late-interaction encoders, and query-time pruning baselines. It measures recovery of exhaustive top-K rankings through overlap, coverage, and latency, including stricter and looser efficiency thresholds.

  • Datasets and Models: The study evaluates five BEIR text corpora and four REAL-MM-RAG multimodal corpora, plus a merged multimodal set for CPU benchmarks.
  • Datasets and Models: ColBERTv2 and Jina-ColBERTv2 provide the late-interaction text encoders, while the visual task uses variable-length queries and 729 document tokens per image.
  • Compared Methods: Full-MaxSim is the exhaustive oracle, compared with Doc-Uniform and Ball-carving query-time pruning baselines using the same corpus tensor.
  • Compared Methods: Doc-Uniform samples ⌈γT⌉ cells uniformly per document and predicts top-K documents from revealed-cell sums.
  • Compared Methods: Ball-carving clusters query tokens by cosine threshold τ, replaces clusters with centroids, and reports coverage as k/T.
  • Efficiency Analysis: Table 5 reports coverage and savings needed to recover 90% of Full-MaxSim’s Overlap@K or nDCG@K, complementing the stricter 95% body-table threshold.

B.2 Detailed Efficiency Results per Dataset

Col-Bandit’s efficiency gains persist across individual BEIR and REAL-MM-RAG datasets, while αef and K expose a controllable speedup–fidelity trade-off. Its runtime advantage is driven primarily by skipping MaxSim cells, with gains remaining robust across dimensions and thread settings.

  • Col-Bandit consistently outperforms the baselines on every individual dataset, with magnitude varying by document length and query difficulty.
  • Smaller αef values increase elimination aggressiveness and speedup but reduce Overlap@5, while the deployed αef=0.2 balances the trade-off.The αef=0.1 setting reaches Ov@5 as low as 0.76, whereas αef=0.3 maintains Ov@5 ≥0.98 on every corpus at reduced speedup.
  • At K=100 with αef=0.2, speedup degrades modestly while Ov@100 remains ≥0.93 on every corpus.
  • At αef=1, Col-Bandit achieves Ov@5=1.0 on every text corpus while retaining nontrivial speedups, despite 28–59% coverage.
  • Col-Bandit maintains near-lossless Top-1 performance and outperforms non-adaptive baselines more clearly at lower coverage budgets.
  • Kernel work accounts for at least 96% of profiled runtime, while elimination and rescore together remain below 4%; thread scaling is sub-linear.Observed speedups can depart from the 100/Cov% heuristic because of memory traffic, thread scaling, and fixed rescore overhead.

C Theoretical Validity in Uniform-Sampling Mode (Special Case)

In uniform-without-replacement mode, Col-Bandit uses Bernstein–Serfling confidence intervals to eliminate documents safely and provides a conditional δ-PAC guarantee at αef=1. The guarantee depends on the simplified radius and exact-precision reveal model, while the practical algorithm progressively rescored survivors exactly.

  • Col-Bandit formulates late-interaction reranking as adaptive Top-K identification by revealing MaxSim entries and returning an estimated exhaustive top-K set.
  • At αef=1, Algorithm 1 has a conditional δ-PAC guarantee for any δ∈(0,1), provided Eq. 7 is a valid per-row confidence radius under uniform-without-replacement reveals.
  • The guarantee applies to a simplified radius with c=1 and omitted O(1/n) terms, and the deployed int8 elimination path receives no PAC claim.The theorem is stated for the exact-precision reveal model; αef=0.2 is treated as a calibrated relaxation.
  • Empirical Bernstein–Serfling bounds provide time-uniform confidence intervals for finite-population row sums under uniform sampling without replacement.A union bound over documents and reveal counts preserves validity at the data-dependent stopping time.
  • A document is safely eliminated when its upper confidence bound falls below the K-th largest lower confidence bound, so repeated elimination does not discard a true top-K document on the high-probability event.
  • The loop stops when at most K+M candidates remain or all tokens are revealed, then exactly rescored survivors determine the returned top-K.

D.2 The NUMKONG C Extension

The NUMKONG C extension implements Col-Bandit as a SIMD-friendly, fused query-time kernel across x86 and ARM. It uses a quantized elimination pass and bit-identical fp32 rescoring, producing substantial measured speedups while preserving ranking fidelity.

  • NUMKONG provides apples-to-apples Full-MaxSim and Col-Bandit comparisons because both use the same C extension and differ mainly in reveal schedule.
  • The implementation supports AVX2 and AVX-512 on x86 and NEON on Apple Silicon, with only the dot-product micro-kernel selected per instruction set.
  • Col-Bandit uses an int8 quantized coarse pass for elimination and an aligned fp32 pass for K+M survivors whose scores are bit-identical to Full-MaxSim.
  • The fused implementation reveals B=4 query tokens per round in a vectorized pass, matching SIMD register tiling and cache-friendly document layouts.
  • Against maxsim-cpu, Col-Bandit is ∼11× faster on an x86 single-thread cross-ISA check, while increasing M from 0 to 5 restores Ov@5 ≥0.98 at small latency cost.

E.1 Extended Related Work

Col-Bandit extends efficient retrieval through query-time, atomic MaxSim-cell pruning rather than index-time compression or token reduction. Its finite-population formulation combines bound-based elimination with Bernstein–Serfling concentration for query-adaptive Top-K identification.

  • Index-time approaches compress representations or prune tokens before seeing the query, whereas Col-Bandit can rerank candidate sets produced by those systems.
  • Col-Bandit is classified as a query-time method that dynamically prunes the atomic interaction matrix H.
  • Col-Bandit applies statistical bound-based early stopping to dense late interaction, pruning atomic MaxSim operations rather than inverted-list pointers.
  • The method estimates the exhaustive top-K from a partially revealed MaxSim matrix, using the actual query to drive elimination.
  • Unlike stochastic-arm identification, each MaxSim row is a finite population of deterministic values sampled without replacement, motivating Bernstein–Serfling concentration.
Loading 2602.02827v2…