Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

Sentence Transformers is a Python library capable of training and employing embedding and reranking models for various applications, including retrieval augmented generation, semantic search, semantic textual similarity, and more. Its latest release, 6.0, features a new MultiVectorEncoder model type, used in ColBERT-style late interaction retrieval, along with a comprehensive training procedure for this encoder.

This article demonstrates how to fine-tune a multi-vector model to surpass general-purpose retrievers on specific datasets. Finetuning multi-vector models requires considering several components, such as the model itself, datasets, loss functions, training arguments, evaluators, and the trainer class. The article explores each of these components and provides practical examples for fine-tuning robust multi-vector models.

Evaluation of the finetuned multi-vector-encoder/mLateOn-medical model, trained in 14.5 hours on a single RTX 3090, reveals its superior performance over general-purpose retrieval models on medical retrieval tasks, such as dense, sparse, lexical, and multi-vector data. For those interested in fine-tuning other types of embedding models, such as dense, sparse, or rerankers, the author recommends reviewing their previous blogposts on training and fine-tuning embedding models, sparse embedding models, and reranker models.

This article primarily focuses on training multi-vector models. For those seeking to learn how to use multi-vector models, from loading and encoding to indexing in vector databases, the companion blogpost on Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers discusses these topics in greater detail. Dense embedding models condense entire text into a single vector, with similarity determined by the product of two such vectors.

In contrast, multi-vector models, also known as late-interaction or ColBERT-style models, exclude this compression and retain a small vector per token. The MaxSim operator scores a query against a document by matching each query token with its best-matching document token, summing the scores for a final similarity measure. These token-level matches preserve fine-grained signals that single vector models tend to average away, leading to stronger retrieval performance with relatively small amounts of in-domain finetuning data.

Additionally, most released retrieval models are configured for short passages, whereas multi-vector models excel at handling longer documents by preserving their full context. When training a custom model, the document length must be configured according to the specific data requirements. For example, a medical evaluation with passages averaging 941 tokens showed that truncation costs up to 0.24 NDCG@10, significantly more than the difference between different model architectures.

Training a MultiVectorEncoder model involves several key components, including the choice of starting point, which can significantly impact the training process. Users can either fine-tune an existing multi-vector model or start from scratch with a new architecture, depending on their data and requirements. A crucial consideration is the length configuration, as many released checkpoints cap document lengths between 180 to 512 tokens.

The mLateOn family accommodates the backbone's full 8192-token context, but if the starting checkpoint has length caps, these should be lifted to avoid truncation during training. The article also highlights a punctuation skiplist that excludes punctuation tokens from document-side scoring and storage, which modestly improves quality while reducing the document index size by 9.6%.

Other modifications, such as using any base transformer as a starting point and allowing training of the token-level projection, further enhance the flexibility and performance of multi-vector models.