Training and Finetuning Multi-Vector Embedding Models

What happened

Sentence Transformers is a Python library for using and training embedding and reranker models for a wide range of applications, such as retrieval augmented generation, semantic search, semantic textual similarity, and more. 0 update introduces a fourth model type: MultiVectorEncoder, for ColBERT-style late interaction retrieval, alongside a complete training approach for it.

The classic ColBERT checkpoints truncate documents at 180 or 300 tokens, and many popular dense models at 256 or 512, because their MS MARCO-style training data rarely goes beyond that. If your documents are long, these models silently discard most of every document before scoring it. 24 NDCG@10, considerably more than any difference between model architectures.

When you train your own model, you configure the document length that your data needs. LightOn ran into this same dynamic with code retrieval, where general LateOn wasn't enough and they trained LateOn-Code. Your domain, whether that's medical, legal, financial, or your company's internal documents, is not getting an official model. This blogpost shows you how to build it yourself, in a matter of hours, on a single consumer GPU.

Let's take a closer look at each component. Multi-vector training gives you a real choice of starting point, and it matters more than you might expect. If you want to further finetune an existing multi-vector model, you don't have to worry about the architecture at all: The checkpoint brings its own recipe along: its query and document marker tokens, its projection head, its scoring skiplist.

For finetuning, you generally want to keep all of that and change only what your data demands. ), and my medical passages run to 1,400 tokens. The mLateOn family already serves the backbone's full 8192 token context, but if your starting checkpoint carries caps, lift them: With the per-task caps unset, truncation falls back to the tokenizer's model_max_length, which is why I configure that limit at load time above.

Why it matters

In this blogpost, I'll show you how to use it to finetune a multi-vector model that outperforms general-purpose retrievers on your data. This method can also train strong new multi-vector models from scratch. Everything below runs on pip install -U "sentence-transformers[train]". Finetuning multi-vector models involves several components: the model itself, datasets, loss functions, training arguments, evaluators, and the trainer class.

I'll have a look at each of these components, accompanied by practical examples of how they can be used for finetuning strong multi-vector models. 5 hours on a single RTX 3090 alongside this blogpost, easily outperforms every general-purpose retrieval model I could find on my medical retrieval evaluation: dense, sparse, lexical, and multi-vector alike.

If you're interested in finetuning dense embedding models, sparse embedding models, or rerankers instead, then consider reading through my prior Training and Finetuning Embedding Models, Training and Finetuning Sparse Embedding Models, and Training and Finetuning Reranker Models blogposts. A dense embedding model compresses a whole text into a single vector, and similarity is one dot product between two such summaries.

A multi-vector model (also called a late-interaction or ColBERT-style model) skips that compression. It keeps one small vector per token and scores a query against a document with the MaxSim operator, where every query token finds its best-matching document token and the scores are summed.

Token-level matching preserves exactly the fine-grained signals that a single vector has to average away, which usually means stronger retrieval, at the cost of a bigger index. The companion Multi-Vector Embedding Models blogpost covers the architecture, encoding, scoring, and indexing in detail, so I'll keep this section short and get to the training.

Finetuning multi-vector models significantly improves their retrieval performance on your specific domain: the vocabulary, the query style, and the notion of relevance all differ between web search, legal discovery, code search, and scientific literature review.

Because queries and documents are matched token by token, multi-vector models pick up fine-grained domain signals that single-vector models tend to average away, and they respond very well to even modest amounts of in-domain finetuning data. Beyond that, most released retrieval models were configured for short passages.

What to watch

I made one more change, adding a punctuation skiplist that excludes punctuation tokens from document-side scoring and storage.

6% on this data for free: You can also point MultiVectorEncoder at any base transformer, and a fresh, randomly initialized token-level projection is appended for you: That's the classic ColBERT pipeline: a Transformer producing contextualized token embeddings, a token-level Dense projecting each of them down to 128 dimensions, a MultiVectorMask deciding which tokens count during scoring, and a token-level Normalize.