ColBERT-Att: Late-Interaction Meets Attention for Enhanced Retrieval

Abstract: Vector embeddings from pre-trained LLMs form a core component in Neural Information Retrieval systems across a multitude of knowledge extraction tasks. The paradigm of late interaction, introduced in ColBERT, demonstrates high accuracy along with runtime efficiency. However, the current formulation fails to take into account the attention weights of query and document terms, which intuitively capture the "importance" of similarities between them, that might lead to a better understanding of relevance between the queries and documents. This work proposes ColBERT-Att, to explicitly integrate attention mechanism into the late interaction framework for enhanced retrieval performance. Empirical evaluation of ColBERT-Att depicts improvements in recall accuracy on MS-MARCO as well as on a wide range of BEIR and LoTTE benchmark datasets.

• The paper introduces an enhanced ColBERT architecture that incorporates attention weights into the late interaction scoring function to better capture token-level importance.
• It demonstrates notable improvements in retrieval accuracy on benchmarks such as MS-MARCO, BEIR, and LoTTE, validating the effectiveness of attention-based modulation.
• The work offers a scalable solution with negligible inference overhead and proposes an attention regularizer to handle document length variations.

ColBERT-Att: Integrating Attention into Late Interaction for Enhanced Neural Retrieval

Overview

The paper "ColBERT-Att: Late-Interaction Meets Attention for Enhanced Retrieval" (2603.25248) presents an extension to the late-interaction paradigm for neural information retrieval, building upon the ColBERT architecture. ColBERT-Att introduces explicit incorporation of attention weights from transformer-based encoders within the late interaction scoring function, aiming to capture token-level importance in both queries and candidate documents. The architecture addresses a critical gap in ColBERT's MaxSim operation, which treats all matched tokens equally regardless of their contextual importance. By integrating query and document attention through a modified scoring function, ColBERT-Att achieves improved retrieval accuracy across both in-domain (MS-MARCO) and out-of-domain (BEIR, LoTTE) benchmarks.

Late Interaction and Attention Integration

ColBERT utilizes a multi-vector representation at token-level granularity, computing relevance as the sum of maximum cosine similarities (MaxSim) between query tokens and document tokens. However, the original formulation ignores contextual importance as signaled by attention weights, potentially overvaluing trivial token matches.

ColBERT-Att augments the scoring function:

$$\mathcal{S}_{\mathcal{Q},~\mathcal{D}} = \sum_{i=1}^{n} e^{\mathcal{A}_{q_i}} \cdot \max_{j=1}^m (\mathcal{E}_{q_{i}} \odot \mathcal{E}_{d_{j}}) \cdot (e^{\mathcal{A}_{d_{w}}})^{\delta}$$

where $\mathcal{A}_{q_i}$ and $\mathcal{A}_{d_{w}}$ are attention weights for query token $q_i$ and the most similar document token $d_{w}$. The exponentiation of attention weights propagates their impact, emphasizing semantically salient terms. To compensate for attention weight variance due to document length, a regularizer $\delta$ is introduced, scaling the attention impact for shorter documents.

This mechanism allows ColBERT-Att to dynamically modulate the contribution of individual token similarities, improving discrimination of relevant passages especially when linguistic overlap is superficial.

Empirical Evaluation

ColBERT-Att is evaluated across MS-MARCO, BEIR, and LoTTE datasets. Training employs ColBERTv2$_{\text{PLAID}}$ as a backbone, with attention weights extracted during encoding and incorporated offline for documents. Inference latency is unaffected because attention weights are a byproduct of transformer encoding.

Numerical Results

• On MS-MARCO, ColBERT-Att demonstrated a Recall@100 improvement of 0.18% over ColBERTv2$_{\text{PLAID}}$, reaching 91.54%. Given the already saturated performance regime, this margin is nontrivial.
• On LoTTE (Success@5), ColBERT-Att outperformed baselines across all domains with weighted averages of 73.5% (Search) and 65.1% (Forum), yielding increases of $\sim 1$% over ColBERTv2$_{\text{PLAID}}$.
• On BEIR search and semantic relatedness tasks (nDCG@10), ColBERT-Att achieved up to 2% gains on challenging datasets (e.g., ArguAna and FiQA), and maintained parity with original ColBERTv2$_{\text{PLAID}}$ elsewhere.
• Ablation studies confirm the necessity of integrating both query and document attention weights; omitting either degrades retrieval efficacy. The attention weight regularizer significantly improves robustness to document length variation, with nDCG@10 on Quora increasing by up to 5% following proper regularization.

Theoretical and Practical Implications

By explicitly integrating attention, ColBERT-Att bridges lexical and semantic matching, similar in spirit to hybrid approaches (e.g., BM42 [vasnetsov2024bm42]) but within the late-interaction neural paradigm. The explicit modulation of similarity using contextual importance signals strengthens model robustness, particularly for heterogenous and long-tailed datasets (LoTTE, BEIR).

Practically, ColBERT-Att retains efficient document representation pre-computation and query-time latency advantages of ColBERT, since attention weights incur negligible additional computation. The introduced attention regularizer ($\delta$) ensures generalization across datasets with divergent text lengths, a frequent challenge in cross-domain retrieval evaluation.

Theoretically, the work suggests that neural retrieval architectures should leverage inherent importance signals, such as attention, rather than treating embedding similarities as uniformly meaningful. This provides new vistas for designing scoring functions that more faithfully model user intent and document salience.

Future Outlook

ColBERT-Att lays groundwork for further exploration of late-interaction retrieval architectures. Potential directions include:

• Integration with advanced transformer features (e.g., rotary position embeddings [su2024roformer], ModernColBERT [GTE-ModernColBERT]).
• Investigation of alternative or learned attention modulation functions beyond simple exponentiation.
• Extension to multi-stage retrieval and answer generation pipelines, such as RAG architectures.

Given that ColBERT-Att can be seamlessly incorporated into existing deployment pipelines with minimal overhead, it is poised to advance both research and industrial search systems, particularly in domains requiring nuanced contextual relevance modeling.

Conclusion

ColBERT-Att enhances late-interaction neural retrieval by embedding explicit attention weights in the token similarity aggregation process. Empirical results across MS-MARCO, BEIR, and LoTTE benchmarks substantiate consistent accuracy gains, especially in settings where traditional ColBERT scoring undervalues term importance. The proposed attention regularizer effectively mitigates document length mismatch, advancing retrieval robustness. ColBERT-Att establishes a compelling case for importance-aware, interaction-based scoring functions in large-scale neural retrieval.