Model Card for ColBERT

ColBERT-v2 is a passage retrieval model that leverages late interaction to allow token-level interaction between query and document embeddings. It outperforms traditional retrievers (such as Dense Passage Retrieval) that encode queries and documents into one-dimensional embeddings.

With the approval of the original author, this model is re-implemented with the huggingface-compatible implementation of PreFLMR. All the weights are imported from the official ColBERTv2 checkpoint.

Model Details

Model Description

Developed by: Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, Matei Zaharia
Re-implemented by: Weizhe Lin, Howard Mei, Jinghong Chen, Bill Byrne
Model type: FLMRModelForRetrieval
Language(s) (NLP): English
License: MIT License

Model Sources

Repository: https://github.com/stanford-futuredata/ColBERT
Repository of implementation: https://github.com/LinWeizheDragon/FLMR
Paper: https://arxiv.org/abs/2112.01488

Uses

Direct Use

This model can be used directly to retrieve documents from a large corpus using text queries. The retrieval usage can be found in the PreFLMR implementation.

Downstream Use

This model can be used combined with language models to create a retrieval-augmented language model.

How to Get Started with the Model

For details of training, indexing, and performing retrieval, please refer to here.

Install the FLMR package.
A simple example use of this model:

from flmr import FLMRConfig, FLMRModelForRetrieval, FLMRQueryEncoderTokenizer, FLMRContextEncoderTokenizer
checkpoint_path = "LinWeizheDragon/ColBERT-v2"
query_tokenizer = FLMRQueryEncoderTokenizer.from_pretrained(checkpoint_path, subfolder="query_tokenizer")
context_tokenizer = FLMRContextEncoderTokenizer.from_pretrained(checkpoint_path, subfolder="context_tokenizer")

model = FLMRModelForRetrieval.from_pretrained(checkpoint_path,
                                query_tokenizer=query_tokenizer,
                                context_tokenizer=context_tokenizer,
                                )

Q_encoding = query_tokenizer(["What is the capital of France?", "What is the capital of China?"])
D_encoding = context_tokenizer(["Paris is the capital of France.", "Beijing is the capital of China.",
                            "Paris is the capital of France.", "Beijing is the capital of China."])

inputs = dict(
    query_input_ids=Q_encoding['input_ids'],
    query_attention_mask=Q_encoding['attention_mask'],
    context_input_ids=D_encoding['input_ids'],
    context_attention_mask=D_encoding['attention_mask'],
    use_in_batch_negatives=True,
)

res = model.forward(**inputs)

Training Details

Training Data

The model was trained on MS MARCO Passage Ranking

More details can be found in the official ColBERT repository

Training Procedure

The training details can be found in the official ColBERT repository and the ColBERTv2 paper

Evaluation

Testing Data, Factors & Metrics

Testing Data

The model is evaluated on a range of retrieval tasks, such as:

Results

Find the evaluation results in the ColBERTv2 paper

Citation

BibTeX:

Please cite the original papers:

@inproceedings{10.1145/3397271.3401075,
author = {Khattab, Omar and Zaharia, Matei},
title = {ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT},
year = {2020},
isbn = {9781450380164},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3397271.3401075},
doi = {10.1145/3397271.3401075},
abstract = {Recent progress in Natural Language Understanding (NLU) is driving fast-paced advances in Information Retrieval (IR), largely owed to fine-tuning deep language models (LMs) for document ranking. While remarkably effective, the ranking models based on these LMs increase computational cost by orders of magnitude over prior approaches, particularly as they must feed each query-document pair through a massive neural network to compute a single relevance score. To tackle this, we present ColBERT, a novel ranking model that adapts deep LMs (in particular, BERT) for efficient retrieval. ColBERT introduces a late interaction architecture that independently encodes the query and the document using BERT and then employs a cheap yet powerful interaction step that models their fine-grained similarity. By delaying and yet retaining this fine-granular interaction, ColBERT can leverage the expressiveness of deep LMs while simultaneously gaining the ability to pre-compute document representations offline, considerably speeding up query processing. Crucially, ColBERT's pruning-friendly interaction mechanism enables leveraging vector-similarity indexes for end-to-end retrieval directly from millions of documents. We extensively evaluate ColBERT using two recent passage search datasets. Results show that ColBERT's effectiveness is competitive with existing BERT-based models (and outperforms every non-BERT baseline), while executing two orders-of-magnitude faster and requiring up to four orders-of-magnitude fewer FLOPs per query.},
booktitle = {Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval},
pages = {39–48},
numpages = {10},
keywords = {bert, deep language models, efficiency, neural ir},
location = {Virtual Event, China},
series = {SIGIR '20}
}

@inproceedings{santhanam-etal-2022-colbertv2,
    title = "{C}ol{BERT}v2: Effective and Efficient Retrieval via Lightweight Late Interaction",
    author = "Santhanam, Keshav  and
      Khattab, Omar  and
      Saad-Falcon, Jon  and
      Potts, Christopher  and
      Zaharia, Matei",
    editor = "Carpuat, Marine  and
      de Marneffe, Marie-Catherine  and
      Meza Ruiz, Ivan Vladimir",
    booktitle = "Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies",
    month = jul,
    year = "2022",
    address = "Seattle, United States",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2022.naacl-main.272",
    doi = "10.18653/v1/2022.naacl-main.272",
    pages = "3715--3734",
    abstract = "Neural information retrieval (IR) has greatly advanced search and other knowledge-intensive language tasks. While many neural IR methods encode queries and documents into single-vector representations, late interaction models produce multi-vector representations at the granularity of each token and decompose relevance modeling into scalable token-level computations. This decomposition has been shown to make late interaction more effective, but it inflates the space footprint of these models by an order of magnitude. In this work, we introduce ColBERTv2, a retriever that couples an aggressive residual compression mechanism with a denoised supervision strategy to simultaneously improve the quality and space footprint of late interaction. We evaluate ColBERTv2 across a wide range of benchmarks, establishing state-of-the-art quality within and outside the training domain while reducing the space footprint of late interaction models by 6{--}10x.",
}

@inproceedings{10.1145/3511808.3557325,
author = {Santhanam, Keshav and Khattab, Omar and Potts, Christopher and Zaharia, Matei},
title = {PLAID: An Efficient Engine for Late Interaction Retrieval},
year = {2022},
isbn = {9781450392365},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3511808.3557325},
doi = {10.1145/3511808.3557325},
abstract = {Pre-trained language models are increasingly important components across multiple information retrieval (IR) paradigms. Late interaction, introduced with the ColBERT model and recently refined in ColBERTv2, is a popular paradigm that holds state-of-the-art status across many benchmarks. To dramatically speed up the search latency of late interaction, we introduce the Performance-optimized Late Interaction Driver (PLAID) engine. Without impacting quality, PLAID swiftly eliminates low-scoring passages using a novel centroid interaction mechanism that treats every passage as a lightweight bag of centroids. PLAID uses centroid interaction as well as centroid pruning, a mechanism for sparsifying the bag of centroids, within a highly-optimized engine to reduce late interaction search latency by up to 7x on a GPU and 45x on a CPU against vanilla ColBERTv2, while continuing to deliver state-of-the-art retrieval quality. This allows the PLAID engine with ColBERTv2 to achieve latency of tens of milliseconds on a GPU and tens or just few hundreds of milliseconds on a CPU at large scale, even at the largest scales we evaluate with 140M passages.},
booktitle = {Proceedings of the 31st ACM International Conference on Information \& Knowledge Management},
pages = {1747–1756},
numpages = {10},
keywords = {colbert, dynamic pruning, efficient search, late interaction, neural information retrieval},
location = {Atlanta, GA, USA},
series = {CIKM '22}
}

If the implementation helps your research, please cite:

@article{Lin_Mei_Chen_Byrne_2024, 
        title={PreFLMR: Scaling Up Fine-Grained Late-Interaction Multi-modal Retrievers}, 
        url={http://arxiv.org/abs/2402.08327}, 
        number={arXiv:2402.08327}, 
        publisher={arXiv}, 
        author={Lin, Weizhe and Mei, Jingbiao and Chen, Jinghong and Byrne, Bill}, 
        year={2024}}