Incorporating Token Importance in Multi-Vector Retrieval

Authors

  • Archish S Microsoft Research
  • Ankit Garg Microsoft Research
  • Kirankumar Shiragur Microsoft Research
  • Neeraj Kayal Microsoft Research

DOI:

https://doi.org/10.1609/aaai.v40i39.40566

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.

Downloads

Published

2026-03-14

How to Cite

S, A., Garg, A., Shiragur, K., & Kayal, N. (2026). Incorporating Token Importance in Multi-Vector Retrieval. Proceedings of the AAAI Conference on Artificial Intelligence, 40(39), 32860–32866. https://doi.org/10.1609/aaai.v40i39.40566

Issue

Section

AAAI Technical Track on Natural Language Processing IV