# ColBERT: Improving Retrieval Performance with Token Level Vector Embeddings

- Canonical: https://33rdsquare.com/colbert-improve-retrieval-performance-with-token-level-vector-embeddings/
- Published: 2024-09-03
- Author: Jordan Brown
- Categories: [Artificial Intelligence & Machine Learning & ChatGPT](https://33rdsquare.com/category/tech/ai/)

---

## Introduction

Retrieval augmented generation (RAG) has emerged as a powerful technique for endowing large language models (LLMs) with the ability to generate more factual and grounded responses. By retrieving relevant context from a knowledge corpus and conditioning the LLM‘s output on that context, RAG helps mitigate the issue of hallucinations – where LLMs confidently generate statements that are incorrect or inconsistent with factual knowledge [1].

At the core of RAG is an information retrieval (IR) component that encodes queries and documents into a shared vector space, allowing for efficient similarity search. Traditionally, bi-encoder models like DPR [2] and BERT [3] have been used for this task, which independently map queries and documents to single embedding vectors. While computationally efficient, compressing an entire text sequence into a fixed-length vector creates an information bottleneck that can limit retrieval performance, especially for more complex queries and documents.

Enter ColBERT – a novel bi-encoder architecture introduced by Khattab and Zaharia in 2020 [4] that represents text as a set of token-level embedding vectors rather than a single vector. By allowing for richer interactions between query and document tokens, ColBERT is able to better capture fine-grained semantic similarities and achieve state-of-the-art retrieval performance on benchmark datasets. In this blog post, we‘ll take a deep dive into the technical details behind how ColBERT works, examine the empirical results, and discuss the broader implications for RAG and other retrieval-based applications.

## The Limits of Single-Vector Text Embeddings

First, let‘s examine why compressing an entire text sequence into a single embedding vector can be problematic for retrieval. Consider the following passage:

"The Amazon rainforest, covering much of northwestern Brazil and extending into Colombia, Peru and other South American countries, is the world‘s largest tropical rainforest, famed for its biodiversity. It‘s crisscrossed by thousands of rivers, including the powerful Amazon. River towns, with 19th-century architecture from rubber-boom days, include Brazil‘s Manaus and Belém and Peru‘s Iquitos and Puerto Maldonado."

A standard bi-encoder model would map this 74-word paragraph into a single D-dimensional embedding vector, which might look something like:

[0.23, -0.12, 0.55, …, -0.37]

where D is typically in the range of 768-1024 for models like BERT.

No matter how large we make D, there will inevitably be some loss of granular information when squeezing everything into a single vector. Nuanced details about the location, characteristics, cities, and history of the Amazon get blurred together.

This isn‘t a huge issue for simple keyword searches, but becomes problematic when dealing with more complex queries. Imagine a user asks:

"What are some major cities located along the Amazon river?"

To accurately retrieve the above passage, the model needs to understand the compositional relationship between "major cities", "located", and "Amazon river". Decoding that from a single vector is challenging. The embedding might capture that the text is broadly about the Amazon and cities, but may lose the specific detail about those cities being located on the river.

Empirically, Khattab and Zaharia showed that single-vector models often fail to retrieve semantically relevant documents that don‘t contain exact keyword matches to the query. On the Natural Questions benchmark, which consists of real queries from Google Search, a BERT bi-encoder was only able to retrieve a relevant passage for 59% of queries, compared to 74% for BM25, a classic keyword-based retrieval method.

## ColBERT‘s Multi-Vector Approach

ColBERT addresses this limitation by representing text as a set of token-level embedding vectors, allowing it to capture more granular relationships between queries and documents.

Instead of a single D-dimensional vector for the full text sequence, ColBERT produces a L x D embedding matrix, where L is the number of tokens in the sequence. Each row of this matrix is a D-dimensional embedding for one token.

For our Amazon example paragraph, the ColBERT embeddings would look something like:

"The": [0.12, 0.52, …]
 "Amazon": [-0.21, 0.34, …]
 "rainforest": [0.63, -0.09, …]
 …
 "Maldonado": [-0.47, 0.28, …]

At query time, ColBERT embeds the query tokens in the same way:

"major": [0.54, 0.03, …]
 "cities": [-0.32, 0.19, …]
 "Amazon": [0.18, -0.61, …]
 "river": [0.42, 0.37, …]

Then, instead of comparing a single query vector to a single document vector, ColBERT performs a "late interaction" step. It computes the cosine similarity between each query token and each document token, producing a L_Q x L_D similarity matrix.

In our cities example, this matrix would likely show high similarity scores between:

- "major" and "Manaus", "Belém", "Iquitos", "Puerto"
- "cities" and "towns"
- "Amazon" and "Amazon"
- "river" and "rivers", "Amazon"

ColBERT then computes a "MaxSim" relevance score by taking the maximum similarity for each query token and summing those maximum values. Intuitively, this corresponds to matching each query token to its most similar counterpart in the document.

Formally, for a query Q and document D, the ColBERT relevance score is defined as:

$$ \text{Score}(Q, D) = \sum_{i=1}^{L_Q} \max_{j=1}^{L_D} \text{cos_sim}(Q_i, D_j) $$

Through this late interaction mechanism, ColBERT is able to identify that the given passage does indeed contain relevant information to the query, even though the exact phrase "major cities" is not present. The multi-vector representation provides a more nuanced understanding of the semantic relationships between the key entities and concepts.

This approach has several advantages over single-vector models:

1. It allows for more precise matching between queries and documents by capturing fine-grained similarities at the token level. This is especially important for complex, compositional queries that can‘t be answered by simple keyword matching.
2. It reduces the information bottleneck by allowing each token to be represented individually rather than compressed into a single vector. This lets ColBERT match queries to relevant documents even when there isn‘t high overlap between the query and document vocabularies.
3. It enables efficient retrieval by precomputing and indexing the document token embeddings offline. At query time, only the query token embeddings need to be computed, and the MaxSim scores can be calculated using fast matrix multiplication on a GPU.

In their experiments, Khattab and Zaharia found that ColBERT significantly outperformed single-vector models like BERT and DPR on a range of benchmark datasets. On the Natural Questions task, ColBERT achieved a top-5 retrieval accuracy of 82%, a 13% absolute improvement over the BERT bi-encoder. ColBERT also demonstrated strong performance on the MS MARCO passage ranking leaderboard [5], where it held the state-of-the-art at the time of publication.

## Scaling with Compression

One potential downside of using multiple embedding vectors per text sequence is the increased storage requirements compared to single-vector models. ColBERT addresses this in two ways:

1. Dimensionality reduction: ColBERT uses a smaller embedding dimensionality (typically 128) compared to models like BERT which usually have 768-1024 dimensions. This is possible because each individual token embedding doesn‘t need to encode the entire semantic meaning of the sequence.
2. Residual quantization: Only the first token embedding is stored in its full 32-bit float format. The rest of the embeddings are compressed by quantizing them to 8-bit integer values representing the differences from the first embedding. This capitalizes on the fact that token embeddings from the same sequence tend to have similar magnitudes and directions.

Together, these techniques allow ColBERT to store the token embeddings in a highly compressed format while still preserving the benefits of late interaction. In practice, the compressed ColBERT indexes take up 2-4x the space of a single-vector index, but are still manageable for billion-scale document collections. Compression ratios of 20-40x can be achieved relative to storing the raw BERT embeddings.

## Real-World Impact

The strong empirical results and scalability of ColBERT have made it an attractive choice for a range of real-world applications. Some notable examples include:

- Microsoft Bing: In 2020, Microsoft announced that they had integrated ColBERT into the Bing search engine to improve the quality of text snippets shown in search results [6]. By retrieving passages that are more semantically relevant to user queries, ColBERT has helped make Bing‘s search results more useful and engaging.
- Amazon Alexa: Amazon has used ColBERT to build a retrieval system for answering questions about natural language commands for the Alexa virtual assistant [7]. When a user issues a command like "Alexa, play my running playlist", ColBERT is used to find relevant documentation that explains how to execute the command.
- Semantic Scholar: The academic search engine Semantic Scholar uses ColBERT to power their scientific literature recommendation system [8]. By capturing fine-grained semantic similarities between papers, ColBERT is able to surface highly relevant papers even when they don‘t share exact keyword matches with the query.

Beyond these specific use cases, the general approach of late interaction has inspired many subsequent works in the field of neural IR. Google‘s SMITH model [9] extends ColBERT‘s ideas to longer text sequences by using a sliding window approach. Facebook‘s RocketQA[10] model combines dense and sparse retrieval in a multi-stage late interaction architecture. Many domain-specific applications like bug localization [11], clinical note retrieval [12], and source code search [13] have also benefited from token-level interaction models.

## Looking Ahead

As the field of neural IR continues to evolve at a rapid pace, it‘s worth reflecting on how approaches like ColBERT fit into the bigger picture and where things may be headed.

One important trend has been the shift towards even more fine-grained interaction mechanisms that go beyond the token level. For example, the Poly-Encoder [14] and ME-BERT [15] models compute attention scores between query and document tokens, allowing for more expressive matching than ColBERT‘s MaxSim operator. The DeepImpact [16] model considers term importance via retrieval-enhanced contextualization.

Another exciting direction is the integration of retrieval into the pre-training process itself. Recent models like SPAR [17] and Atlas [18] augment language model pre-training with a dense retrieval objective, allowing the model to learn token representations that are directly optimized for semantic search. This alleviates the need for a separate post-hoc retrieval step and has led to impressive gains on benchmark datasets.

More broadly, as retrieval systems continue to get better at surfacing relevant information, the line between IR and question answering is starting to blur. Models like REALM [19], RAG [1], and DPR [2] have shown that retrieval can be tightly integrated with LLMs to enable open-domain QA with less hallucination. The goal is to build systems that can engage in open-ended dialogue while still being grounded in factual knowledge. Approaches like ColBERT that improve the quality of the retrieval process are a key enabler for this ambitious vision.

Looking ahead, we can expect to see retrieval systems that are even more sophisticated in their ability to find relevant information and reason over it. Better retrieval has the potential to improve the accuracy, efficiency, and interpretability of NLP systems across a wide range of applications, from search and question answering to text generation and task-oriented dialogue. As a field, it‘s an exciting time to be pushing the boundaries of what‘s possible.

## References

[1] Lewis et al. "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks." arXiv 2005.11401 (2020).
 [2] Karpukhin et al. "Dense Passage Retrieval for Open-Domain Question Answering." EMNLP 2020.
 [3] Devlin et al. "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding." NAACL 2019.
 [4] Khattab and Zaharia. "ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT." SIGIR 2020.
 [5] MS MARCO Passage Ranking Leaderboard. [https://microsoft.github.io/msmarco/](https://microsoft.github.io/msmarco/)
 [6] "Bing delivers its largest improvement in search experience using Azure GPUs and NVIDIA TensorRT" Microsoft Research Blog, 2020.
 [7] Kullish et al. "Meta Retrieval for User Queries with Compositional Intent" CIKM, 2022.
 [8] Lauscher et al. "Specializing Document Embeddings for Scientific Domains using Fine-tuned ColBERT." arXiv 2201.01772 (2022).
 [9] Yang et al. "Beyond 512 Tokens: Siamese Multi-depth Transformer-based Hierarchical Encoder for Long-Form Document Matching." CIKM 2020.
 [10] Qu et al. "RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering." NAACL 2021.
 [11] Huang et al. "BIKER: Bug Identification using Knowledge-enhanced Retrieval." ASE 2022.
 [12] Di Nunzio et al. "Leveraging Language Models and Late Interaction for Clinical Semantic Textual Similarity and Natural Language Inference." IberLEF@SEPLN 2022.
 [13] Zhou et al. "Code Search Intent Classification Using ColBERT Late Interaction over Hybrid Representations." arXiv 2111.02229 (2021).
 [14] Humeau et al. "Poly-encoders: Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring." ICLR 2020.
 [15] Luan et al. "Sparse, Dense, and Attentional Representations for Text Retrieval." TACL 2021.
 [16] Mallia et al. "DeepImpact: Toward Effective and Efficient IR with Contextual Term Weighting." SIGIR 2021.
 [17] Chen et al. "SPAR: Attention-based Sparse Retriever for Large-scale Passage Retrieval." arXiv 2204.07541 (2022).
 [18] Izacard et al. "Atlas: Few-shot Learning with Retrieval Augmented Language Models." ICLR 2022.
 [19] Guu et al. "REALM: Retrieval-Augmented Language Model Pre-Training." ICML 2020.

---

Source: [ColBERT: Improving Retrieval Performance with Token Level Vector Embeddings](https://33rdsquare.com/colbert-improve-retrieval-performance-with-token-level-vector-embeddings/)