Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

Summary

Multi-vector embedding models, also known as late-interaction or ColBERT-style architectures, avoid compressing an entire text into one vector. Instead they retain a contextualized embedding per token and score queries against documents with the MaxSim operator, which sums the best token-level matches. This design preserves fine-grained signals that single-vector models must average away, at the expense of larger indexes and higher memory use during retrieval.

Finetuning these models on domain-specific data yields clear gains because vocabulary, query style and relevance criteria vary sharply across fields. On long documents the benefit is especially pronounced: many released checkpoints truncate passages at 180–512 tokens, silently discarding content that averages 941 tokens in medical collections and costing up to 0.24 NDCG@10. Training on in-domain pairs lets practitioners set document length to match their data and avoids the performance regression observed when continuing from fully supervised general-purpose checkpoints.

The training pipeline in Sentence Transformers centers on the MultiVectorEncoder class together with MultiVectorMultipleNegativesRankingLoss (or its memory-efficient GradCache variant CachedMultiVectorMultipleNegativesRankingLoss). Datasets are supplied as simple (query, relevant passage) pairs; larger effective batch sizes improve contrastive learning while the cache mechanism keeps GPU memory bounded. Training arguments control learning rate, batch sizing and length limits, and the MultiVectorInformationRetrievalEvaluator supplies concrete NDCG@10 feedback on held-out domain data.

Ablation experiments across six starting checkpoints on 25 k medical pairs showed that unsupervised retrieval-pretrained models adapt more readily than their supervised siblings, which often stagnate or regress. The resulting domain-adapted model, trained for 14.5 hours on a single RTX 3090, surpassed every dense, sparse, lexical and multi-vector baseline evaluated on the medical test set, delivering a +0.062 NDCG@10 lift while fitting within 17.5 GB VRAM. Additional index-size reductions are obtained through pooling and quantization without further quality loss.

Why it matters

Directly actionable for ML Engineers: provides code, loss scaling guidance, document-length handling, and index optimization that teams can apply immediately for domain-specific retrieval on long documents common in Dutch healthcare, legal, and enterprise use cases.