Fine-tuning ColBERT-style multi-vector models with Sentence Transformers v6.0 — a medical retrieval case

Sentence Transformers is a Python library for using and training embedding and reranker models across tasks like retrieval-augmented generation, semantic search, and semantic textual similarity. Version 6.0 introduces a new model type, MultiVectorEncoder, for ColBERT-style late-interaction retrieval and ships a complete training approach for it. Examples in the post run after pip install -U "sentence-transformers[train]".

The author demonstrates how to finetune or train multi-vector models from scratch by assembling model, dataset, loss, training arguments, evaluators and the trainer. As a concrete case, the author trained a model called multi-vector-encoder/mLateOn-medical in 14.5 hours on a single RTX 3090 (peak ~17.5 GB VRAM) and reports it outperforms a broad set of general-purpose retrievers on an in-domain medical benchmark.

What is a Multi-Vector model?

A dense embedding model compresses an entire text into a single vector and measures similarity with one dot product. A multi-vector (late-interaction or ColBERT-style) model keeps one small vector per token and scores queries with a MaxSim operator: each query token finds its best-matching document token and those scores are summed. Token-level matching preserves fine-grained signals that a single-vector representation tends to average away, typically yielding stronger retrieval at the cost of a larger index.

The companion post on Multi-Vector Embedding Models covers architecture and indexing in depth; here the focus is on training.

Why finetune?

Finetuning multi-vector models substantially improves retrieval performance on a target domain because vocabulary, query style and relevance notions differ between domains (web search, legal, code, scientific or medical literature). Token-level matching makes multi-vector models especially receptive to modest amounts of in-domain finetuning data.

Additionally, many released retrievers were configured for short passages (e.g., classic ColBERT checkpoints truncate at 180 or 300 tokens; many dense models at 256 or 512). If your documents are long (the author’s medical passages average 941 tokens), truncation can cost a lot — the author measured up to 0.24 NDCG@10 lost due to truncation. Training your own model lets you set the document length your data requires.

Training components

Training MultiVectorEncoder models requires these components:

• Model (finetune or build fresh)
• Dataset (training and evaluation)
• Loss function
• Training arguments (performance and logging)
• Evaluator (optional but recommended)
• Trainer (glues everything together)

Below are the practical choices and recommendations from the author.

Model: finetune an existing checkpoint or build from a base transformer

If you finetune an existing multi-vector checkpoint (e.g. lightonai/mLateOn-unsupervised), you generally keep the architecture and checkpoint recipe (query/document markers, projection head, scoring skiplist) and change only what the data requires. Check and lift any query_length or document_length caps if your documents are longer than the checkpoint’s defaults. The author loads checkpoints preferring fp32 when memory allows and sets processor_kwargs model_max_length=8192 for long-medical passages.

The author also applied a punctuation skiplist on the document side (excluding punctuation tokens from scoring/storage), which produced a small quality gain and ~9.6% smaller document index on the data tested.

Alternatively, MultiVectorEncoder can wrap any base transformer and append a fresh token-level projection layer (e.g., projecting contextualized token embeddings to 128 dims). A randomly initialized projection needs training but can reach competitive performance quickly when paired with a strong backbone.

Which starting point works best?

The author compared six starting points trained identically on 25k medical question–passage pairs and evaluated on 1,000 held-out questions searching 50,000 passages. A repeated finding: "-unsupervised" (contrastively pretrained but not supervised-finetuned) checkpoints adapted to the new domain far better than fully finished checkpoints. These pre-supervised checkpoints carry the late-interaction structure without general-purpose supervised tuning that the domain finetune would have to undo. If available, start from a pre-supervised checkpoint; otherwise, a fresh projection on a strong retrieval-pretrained backbone is a close second.

Dataset

MultiVectorEncoderTrainer uses datasets.Dataset or datasets.DatasetDict. You can load data from the Hugging Face Datasets Hub or use local files (CSV, JSON, Parquet, Arrow, SQL). Many public datasets working out of the box are tagged with sentence-transformers on the Hub.

The example dataset is tomaarsen/miriad-4.4M-split (about 4.4M medical question–passage pairs, passages average 941 tokens). For the long run the author selected 1,000,000 pairs from that dataset for training.

Dataset format must match the loss function. If a loss requires a label column, your dataset needs a "label" or "score" column. For multi-vector tasks, the first column is treated as the query and subsequent columns as documents unless overridden by router_mapping. Knowledge-distillation datasets can be laid out as (query, document_1, ..., document_N, scores) where scores are teacher scores.

Loss function

For question–passage pairs the standard approach is in-batch negatives with MultiVectorMultipleNegativesRankingLoss. The author recommends CachedMultiVectorMultipleNegativesRankingLoss (GradCache variant) to decouple effective contrastive batch size from GPU memory by chunking documents; mini_batch_size controls chunk size (the author uses mini_batch_size=16 for 940-token documents, equivalent to mini_batch_num_tokens15,000).

Crucially, contrastive losses for multi-vector models default to scale=1.0 (not 20.0 as in dense training), because MaxSim sums per-token similarities and already spans a larger range; copying dense-scale=20.0 would saturate the softmax and harm gradients.

For distillation from a stronger teacher, use MultiVectorDistillKLDivLoss.

Training arguments

Use MultiVectorEncoderTrainingArguments to control training. The author’s actual run used these notable settings:

• output_dir: models/mLateOn-medical
• num_train_epochs: 1
• per_device_train_batch_size: 128 (effective contrastive batch thanks to GradCache)
• per_device_eval_batch_size: 16
• learning_rate: 1e-4
• warmup_steps: 0.05
• prompts mapping: {"question": "[Q] ", "passage_text": "[D] "}
• fp16: False, bf16: True
• batch_sampler: BatchSamplers.NO_DUPLICATES
• eval/save/log strategies and steps configured for frequent checkpoints/logging

Notes: prompts must be mapped explicitly for training; max_length should be left unset to match inference unless you accept the quality/speed trade-off; the author found learning_rate=1e-4 worked best after a sweep.

Evaluator

Concrete retrieval metrics are most informative. Sentence Transformers provides multi-vector evaluators including MultiVectorInformationRetrievalEvaluator; for domain finetuning the author used that with held-out queries, a corpus and relevant-doc mappings. To avoid saturated evals (e.g., gold passages too easy), add distractor passages until scores spread out — the author built a corpus with about 200k passages by adding deduplicated train passages as distractors.

Trainer and the full training recipe

The MultiVectorEncoderTrainer unifies the components. The author’s full recipe that produced multi-vector-encoder/mLateOn-medical:

1. Load lightonai/mLateOn-unsupervised in fp32, processor_kwargs model_max_length=8192 and set model_card_data.
2. Unset per-task caps: model[0].query_length=None and model[0].document_length=None.
3. Add punctuation skiplist to model[2] and resolve with tokenizer.
4. Load 1,000,000 training pairs from tomaarsen/miriad-4.4M-split.
5. Use CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=16).
6. Small dev evaluator: 500 held-out questions vs ~10k unique passages to watch progress.
7. Training arguments as above (per_device_train_batch_size=128, learning_rate=1e-4, bf16 True, etc.).
8. trainer.train(); save final model and optionally push to Hugging Face Hub.

The run took 14.5 hours on a single RTX 3090 with peak VRAM ~17.5 GB. Scaling experiments indicate 100k pairs (~75 minutes) was within ~0.012 NDCG@10 of the full 1M run; most gains arrive early.

Callbacks and multi-dataset training

The trainer supports transformers.TrainerCallback subclasses like WandbCallback, TensorBoardCallback and CodeCarbonCallback, enabled via report_to. For multi-dataset training, pass a dictionary of datasets and optionally a dictionary of losses; sampling strategies include ROUND_ROBIN or PROPORTIONAL.

Evaluation results on the MIRIAD benchmark

The author evaluated the finetuned model against over 50 retrieval configurations across architecture families on a 1,000-query / 200,000-passage MIRIAD benchmark (10k gold passages hidden among 190k deduplicated distractors). Main headline:

• multi-vector-encoder/mLateOn-medical (finetuned): NDCG@10 = 0.9139, acc@1 = 0.849
• lightonai/mLateOn (zero-shot): NDCG@10 = 0.8520, acc@1 = 0.758
• lightonai/GTE-ModernColBERT-v1 (zero-shot, cap lifted): NDCG@10 = 0.8502
• Qwen/Qwen3-Embedding-4B (dense, zero-shot): NDCG@10 = 0.7817
• BM25 (lexical): NDCG@10 = 0.7501

The finetuned model beats the best zero-shot model of any architecture by +0.062 NDCG@10. In practical terms, the best zero-shot model returns the correct passage as rank 1 for 75.8% of queries, while the finetuned model does so for 84.9%, reducing rank-1 error by over one third on this dataset.

The architecture trend is clear: late-interaction (multi-vector) models top the table on long documents. Scale alone does not rescue single-vector models; Qwen3-Embedding-4B (an otherwise strong dense model) still trails by ~0.13 NDCG@10.

Caveat: MIRIAD’s queries are generated from the passages, so lexical baselines like BM25 do well here — but that does not necessarily generalize to other domains.

Optimizing the index: size vs. accuracy

A common objection to multi-vector retrieval is index size. For the author’s data the raw token vectors require ~878 vectors per passage, and a 200k-passage corpus is ~45 GB in fp16. But compression strategies make multi-vector indexes practical:

• HierarchicalTokenPooling clusters token embeddings and stores cluster centroids, keeping roughly 1/pool_factor of vectors. The author measured small accuracy losses for large reductions (e.g., halving vectors cost 0.0033 NDCG@10, quartering still scored 0.8991).

• PLAID-style quantization plus pruning (measured by Omar Khattab using fast-plaid with 1-bit residuals) produced far smaller indexes: 1-bit PLAID with all vectors produced a 3.37 GB index (13× smaller than raw embeddings) with just 0.0155 NDCG@10 drop. Aggressive pruning could reduce the index to 1.45 GB while still scoring higher than certain large dense models.

Quantization should be the first lever; pooling and pruning further reduce storage. Properly configured indexes make the index-size objection much less compelling.

Conclusions

Sentence Transformers v6.0 brings first-class support for ColBERT-style late-interaction models via MultiVectorEncoder and an integrated training pipeline. Domain finetuning — in the author’s recipe a pre-supervised checkpoint, one million domain pairs, in-batch negatives with GradCache, full document length and a relatively high learning rate — produced a model in 14.5 hours on a single RTX 3090 that outperformed a wide range of general-purpose retrievers on a medical retrieval benchmark.

The article gives practical guidance on model choices, dataset formatting, loss functions, training arguments, evaluators, the trainer script, callbacks, multi-dataset training and index optimization strategies (pooling, quantization, pruning). Smaller-scale runs (e.g., 100k pairs) still deliver most of the benefit quickly, making domain finetuning accessible on a single consumer GPU.

Acknowledgements and further reading

Thanks to Omar Khattab for measurements of quantized and pruned index configurations. The post points to additional resources and training examples (MIRIAD, MS MARCO, multimodal, PEFT adapters) and the Sentence Transformers documentation (Installation, Quickstart, Training Overview, Loss Overview, API Reference, Distributed Training) as well as the companion post on using multi-vector models (encoding, indexing and runtime details).