Incorporating Token Importance in Multi-Vector Retrieval
Archish S, Ankit Garg, Kirankumar Shiragur, Neeraj Kayal
TL;DR
This paper extends ColBERT's multi-vector retrieval by introducing Weighted Chamfer, a token-weighted variant of the Chamfer distance used in late interaction scoring. It formalizes the approach with $ ext{Weighted Chamfer} = \eta_w(q,d) = \frac{1}{|q|} \sum_{i=1}^{|q|} w_{q_i} \cdot \min_{j} \|E_q(q_i) - E_d(d_j)\|_2 $, where token weights $w_{q_i}$ can be static (IDF-based) or learned through a contrastive, two-negative-set loss that remains convex in $w$. The authors demonstrate zero-shot gains on BEIR using IDF weights and considerable few-shot gains by training the weights with limited relevance data, all while keeping the base ColBERTv2 encoders fixed and incurring no additional inference latency. They provide theoretical results on the convexity of the training objective and sample-complexity bounds, and show that the approach is storage-efficient and adaptable to low-resource settings. Overall, Weighted Chamfer offers a lightweight, effective enhancement to multi-vector retrieval, improving cross-domain generalization and re-ranking quality without modifying the underlying model architecture.
Abstract
ColBERT introduced a late interaction mechanism that independently encodes queries and documents using BERT, and computes similarity via fine-grained interactions over token-level vector representations. This design enables expressive matching while allowing efficient computation of scores, as the multi-vector document representations could be pre-computed offline. ColBERT models distance using a Chamfer-style function: for each query token, it selects the closest document token and sums these distances across all query tokens. In our work, we explore enhancements to the Chamfer distance function by computing a weighted sum over query token contributions, where weights reflect the token importance. Empirically, we show that this simple extension, requiring only token-weight training while keeping the multi-vector representations fixed, further enhances the expressiveness of late interaction multi-vector mechanism. In particular, on the BEIR benchmark, our method achieves an average improvement of 1.28\% in Recall@10 in the zero-shot setting using IDF-based weights, and 3.66\% through few-shot fine-tuning.