PyLate: Flexible Training and Retrieval for ColBERT Models

ColBERT models are strong out-of-domain retrievers

An important part of retrieval-augmented generation tools such as Paradigm is searching in the user documents collection to find relevant information to answer a query.
A popular architecture to search for relevant documents in large databases is bi-encoders that pre-compute the representation of all the documents to index them. Then, to perform a search, the user query is encoded and the search for the most similar documents is performed by computing the similarity of the query with the pre-computed representation of the documents with a lightweight similarity computation. This setup enables quick search even in large databases.
ColBERT models are another type of encoder also used for retrieval tasks. Unlike traditional bi-encoders that pool all the token representations into a single one, ColBERT models retain all token representations and use late interaction (MaxSim) to compute query/document similarity.

This subtle yet crucial difference leads to strong results and impressive generalization to Out-Of-Domain (OOD) data, i.e, it works well on data different from what the model has been trained on. This is a very important property because it is unrealistic to consider that the training data will cover all the possible domains the model will be used on and poor OOD generalization are the root of a lot of failure cases. Thus, given that the customer data is most likely unique and is usually not leveraged to train the models, having strong OOD performance is crucial to build a pipeline that will work on all the users data.

Introducing PyLate

Compared to the growing interest in late interaction models, there have been few research and released models in this area.

We are releasing PyLate, a library that aims to fill this gap by providing a modular library built on the widely-used sentence-transformers framework, making it both accessible and familiar to many users.

PyLate comes with various features:

1. An interface that is accessible for newcomers and familiar to experienced users
2. Support for any triplet dataset available on Hugging Face for contrastive learning as well as any encoder as the base model
3. Training of the model using knowledge distillation to get the best performance
4. Multi-GPU and FP16/BF16 training with Weights & Biases logging for efficient and trackable training
5. HNSW index to store and query the documents as well as reranking first-stage retrieval pipeline results without index
6. Support of document tokens pooling to reduce the memory footprint of the index without degrading performance

We plan to release various strong late interaction (English and multilingual) models in the near future using PyLate and integrate them in Paradigm to enhance the quality of our retrieval pipeline on our customer data.

If you are interested in building and exploring this promising line of research, you will find documentation and examples on the Github repository: https://github.com/lightonai/pylate

We already have many features planned for the future but we also welcome contributions, so feel free to check the contribution guide and help us build better retrieval models!
‍