Sentence Transformers adds training support for multi-vector ColBERT-style embeddings

Sentence Transformers adds training support for multi-vector ColBERT-style embeddings

A Hugging Face blog post lays out how to train and finetune multi-vector, ColBERT-style embedding models using Sentence Transformers' MultiVectorEncoder class, installable with pip install -U sentence-transformers[train]. Unlike a dense embedding model, which compresses a whole passage into one vector and compares two such summaries with a single dot product, a multi-vector, or late-interaction, model keeps one small vector per token and scores a query against a document with the MaxSim operator: every query token is matched against its best document token, and the scores are summed. That token-level matching preserves fine-grained signal a single vector has to average away. It usually improves retrieval, at the cost of a larger index.

The post gives two reasons to finetune. First, domain adaptation: vocabulary, query style, and the notion of relevance all differ between web search, legal discovery, code search, and scientific literature review. Because multi-vector models match token by token, they pick up fine-grained domain signals easily and respond well to even modest amounts of in-domain data. LightOn hit this directly with code retrieval: its general LateOn model was not enough, so the company trained a dedicated LateOn-Code. Second, document length: classic ColBERT checkpoints truncate documents at 180 or 300 tokens, and many popular dense models cap out at 256 or 512, because their MS MARCO-style training data rarely runs that long. On the author's own medical evaluation, where passages average 941 tokens, that truncation costs up to 0.24 NDCG@10, more than the gap between model architectures.

Training a MultiVectorEncoder model involves six components, each covered in the post: the model itself, a dataset, a loss function that measures performance and guides optimization, optional training arguments that affect training performance, tracking, and debugging, an optional evaluator that scores the model before, during, or after training, and a trainer class that brings the rest together.

On the model, two starting points are available. Continuing an existing multi-vector checkpoint, for example lightonai/mLateOn-unsupervised, keeps its query and document marker tokens, projection head, and scoring skiplist, and needs only attention to length limits. The mLateOn family already serves its backbone's full 8,192 token context, but many other released checkpoints, including GTE-ModernColBERT-v1, which ships with a 48 token query length and a 300 token document length, cap documents well short of the author's 1,400 token medical passages, so those per-task caps get unset in favor of the tokenizer's own limit. A punctuation skiplist that excludes punctuation tokens from document-side scoring and storage modestly improved quality in a four-way ablation testing none, punctuation, stopwords, and both, and it shrank the document index by 9.6% on the author's data for free. Building a model from scratch instead means pointing MultiVectorEncoder at any base transformer, such as answerdotai/ModernBERT-base: a fresh, randomly initialized token-level projection to 128 dimensions is appended automatically, producing the standard ColBERT-style stack of a transformer, a token-level dense projection, a mask deciding which tokens count during scoring, and a token-level normalization step. The classic ColBERT tokenization tricks, mask-token query expansion, query and document prefix tokens, a document-length cap, and a punctuation skiplist, are all off by default. Testing mask-token query expansion in four configurations produced no measurable difference for the author's domain.

To settle which starting point is actually best, the author trained six of them with an identical recipe on 25,000 medical question-passage pairs from MIRIAD and evaluated each on 1,000 held-out questions against a 50,000-passage corpus. The result replicated across two model families: checkpoints still in their unsupervised, pre-supervised-finetuning state adapted to the new domain far better than their fully finetuned siblings, overtaking them despite starting from a lower score. The finished checkpoints, by contrast, barely moved or regressed, at every learning rate tried. A fresh projection on a strong retrieval-pretrained backbone, Alibaba-NLP's gte-modernbert-base, came close: with only the new projection and 25,000 training pairs, it landed within 0.03 of the existing-checkpoint starting points. The post's advice follows directly: start from a pre-supervised checkpoint if the model family publishes one, fall back to a fresh projection on a strong backbone if not, and treat continuing from a fully finished checkpoint as the weakest option despite it feeling the most natural.

On data, the trainer takes a Hugging Face Datasets object, loadable from the Hugging Face Hub or from local CSV, JSON, Parquet, Arrow, or SQL files. The post's own training set is tomaarsen/miriad-4.4M-split: 4.4 million medical question-passage pairs built from MIRIAD, 4,467,542 rows exactly, each pairing a question with the passage that contains its answer. Simple query-passage pairs like these, the post argues, are the easiest retrieval training data to collect for a new domain, and are all that is required.

As a worked example, the author finetuned a model it calls multi-vector-encoder/mLateOn-medical on that MIRIAD data, training it in 14.5 hours on a single RTX 3090, and reports that it easily outperforms every general-purpose retrieval model the author could find on the medical retrieval evaluation, dense, sparse, lexical, and multi-vector alike. No exact score or margin for that comparison, no dollar training cost, and no name for the specific loss function the post recommends for multi-vector training are given in the material available here; those details belong to the post's later Loss Function and Evaluation sections. The finetuned model's weights, though, are released publicly as multi-vector-encoder/mLateOn-medical on the Hugging Face Hub.

Key facts

  • Sentence Transformers' MultiVectorEncoder trains ColBERT-style multi-vector retrieval models, installable with pip install -U sentence-transformers[train].
  • The author's worked example, a model called multi-vector-encoder/mLateOn-medical, was finetuned in 14.5 hours on a single RTX 3090 and is said to outperform every general-purpose retrieval model tested, dense, sparse, lexical, and multi-vector alike, though no exact score is given here.
  • The training set is tomaarsen/miriad-4.4M-split, 4.4 million medical question-passage pairs from MIRIAD with passages averaging 941 tokens; classic ColBERT checkpoints truncate at 180 or 300 tokens and popular dense models at 256 or 512, which can cost up to 0.24 NDCG@10 on long passages.
  • Comparing six starting checkpoints, each trained on 25,000 pairs and evaluated on 1,000 held-out questions against a 50,000-passage corpus, found that unsupervised, pre-supervised-finetuning checkpoints adapt to a new domain far better than fully finetuned ones, which barely moved or regressed.
  • A punctuation skiplist that drops punctuation tokens from scoring modestly improved quality in a four-way ablation and shrank the document index by 9.6%.

Why it matters

Multi-vector, or late-interaction, models, the approach ColBERT popularized, keep one small vector per token instead of compressing a whole passage into a single vector, then score a query against a document by matching each query token to its best document token and summing the results. That preserves signal a single vector has to average away, usually at the cost of a bigger index. Sentence Transformers already had companion guides for training dense, sparse, and reranker models; this post completes the set with MultiVectorEncoder, its tool for training and finetuning ColBERT-style models. It also quantifies a problem most teams do not notice: because classic ColBERT checkpoints truncate documents at 180 or 300 tokens, and many popular dense models at 256 or 512, a corpus of longer documents gets silently cut before it is ever scored, costing up to 0.24 NDCG@10 on the author's own medical evaluation, where passages average 941 tokens.

Who it affects

Anyone building retrieval or retrieval-augmented generation on a specialized corpus where general-purpose models underperform. The post names legal discovery, code search, scientific literature review, medical, financial, and internal-documents search as examples, and points to LightOn's own experience as precedent: its general LateOn model was not enough for code retrieval, so the company trained a dedicated LateOn-Code instead. It also targets practitioners without a large compute budget. The worked example finetunes a full domain model on a single consumer GPU, an RTX 3090, in 14.5 hours.

How to use it

Two starting points are on offer. Continuing an existing multi-vector checkpoint, such as lightonai's mLateOn-unsupervised, needs no architecture work, only attention to length limits, since checkpoints like GTE-ModernColBERT-v1 cap documents at 300 tokens by default. A punctuation skiplist that drops punctuation tokens from document-side scoring modestly improved quality in a four-way ablation and shrank the document index by 9.6% at no cost. Building from scratch instead means pointing MultiVectorEncoder at any base transformer, such as ModernBERT-base; a randomly initialized token-level projection is appended automatically. On which to prefer, the post's own six-way comparison favors a pre-supervised checkpoint where the model family offers one, a fresh projection on a strong retrieval-pretrained backbone as the next-best option, and continuing from a fully finished checkpoint as the weakest choice despite feeling most natural. Training data can come from the Hugging Face Hub or from local CSV, JSON, Parquet, Arrow, or SQL files; the post's own example, 4.4 million medical question-passage pairs from MIRIAD, shows that simple query-passage pairs are all a domain finetune needs.

How solid is it

This is one practitioner's own tutorial and experiment, not a peer-reviewed study. Its headline claim, that the finetuned mLateOn-medical model easily outperforms every general-purpose retrieval model tested, spanning dense, sparse, lexical, and multi-vector approaches, is stated without a precise score or margin in the material available here. The starting-point comparison is more concrete: fixed pair counts, a fixed held-out set, and a result the author reports as replicating across two separate model families, a meaningful check even for a single-author test. The specific loss function the post recommends for multi-vector training is not covered by what could be verified here; it belongs to a later section of the post.

Risks and caveats

The gains described are demonstrated on one domain, medical retrieval via MIRIAD, and one hardware setup. The mLateOn-medical model's weights are released publicly on the Hugging Face Hub, though the post gives no dollar training cost, only the hardware, one RTX 3090, and the time, 14.5 hours. The classic ColBERT tokenization tricks, including mask-token query expansion, are off by default; the author tested mask-token query expansion in four configurations and found no measurable difference, so it is not a given win. And the core tradeoff of the whole approach holds regardless of tuning: keeping a vector per token instead of one per document buys the retrieval quality this post is about, at the cost of a bigger index.

“easily outperforms every general-purpose retrieval model I could find on my medical retrieval evaluation: dense, sparse, lexical, and multi-vector alike.”

— Tom Aarsen, Hugging Face blog