Training and Finetuning Multi-Vector Embedding Models with Sentence Transformers

What it claims

Sentence Transformers now has a MultiVectorEncoder class and trainer for ColBERT-style late-interaction models. The author's main recipe is a pre-supervised checkpoint (lightonai/mLateOn-unsupervised), 1M in-domain query–passage pairs, in-batch negatives, full document length and LR 1e-4. It produced multi-vector-encoder/mLateOn-medical. On MIRIAD retrieval (1,000 queries over 200k passages), that model raised rank-1 hit rate from 75.8% (best zero-shot model) to 84.9%. It used no teacher model and no mined negatives. Most of the post is ablation-backed advice that changes how you would configure such a run.

Evidence

Starting point. Six starting points were trained on 25k MIRIAD pairs and tested against a 50k-passage corpus.

• The -unsupervised checkpoints overtook their finished siblings in two model families, despite starting lower.
• Finished checkpoints barely moved or regressed at every learning rate tried.
• A fresh random projection on Alibaba-NLP/gte-modernbert-base came within 0.03 of the checkpoint starting points.

Document truncation costs more than architecture choice. Passages average 941 tokens. Many released models cap documents at 180–512 tokens.

• Lifting the cap was worth +0.08 to +0.24 NDCG@10 for every multi-vector baseline.
• It was worth +0.03 for DenseOn.

Late interaction vs dense. At matched data and backbone:

• LateOn beats DenseOn by +0.12; mLateOn beats mDenseOn by +0.13.
• Qwen3-Embedding-4B, with about 33x the active parameters, is still 0.13 short of the finetuned model.
• Qwen3-Embedding-8B scores below the 4B version.

Ablations from the training run:

• Training at max_length=512 lost about 0.015 NDCG@10 for about 2x speed. The gap did not close with more data.
• An LR sweep from 5e-6 to 2e-4 favoured 1e-4.
• Batch sizes above 128 gave no further gain.
• [MASK] query expansion made no measurable difference in any of 4 configurations.
• A punctuation skiplist modestly won a 4-way skiplist ablation and cut index size by 9.6%.

Training cost and data scaling:

• The full run took 14.5 h on one RTX 3090 with 17.5 GB peak VRAM.
• 100k pairs (75 min of training) came within 0.012 NDCG@10 of the 1M-pair run.

Index size. The raw index is about 878 vectors per passage, roughly 45 GB at fp16 for 200k passages.

• HierarchicalTokenPooling, applied post-hoc:
• Halving the vector count costs 0.0033 and leaves rank-1 accuracy unchanged.
• Keeping 1/4 of the vectors (11.2 GB) scores 0.8991.
• Keeping 1/10 scores 0.8765.
• Omar Khattab's fast-plaid runs used 1-bit residual quantization, packed ids and document pruning:
• The first configuration was 13x smaller than raw for a 0.0155 loss.
• The most aggressive was 1.45 GB, smaller than Qwen3-Embedding-8B's fp16 index (1.64 GB), and scored 0.0895 higher.

How to use it

pip install -U "sentence-transformers[train]"

Configuration traps the page flags:

• Contrastive loss scale. Contrastive losses default to scale=1.0. Do not copy scale=20.0 from dense scripts: MaxSim scores already span roughly [0, query_length], so 20.0 saturates the softmax and kills gradients.
• Prompts. Training does not apply the model's stored prompts. Map them explicitly, e.g. prompts={"question": "[Q] ", "passage_text": "[D] "}.
• Length caps. Lift per-task caps with model[0].query_length = None and model[0].document_length = None. Then set processor_kwargs={"model_max_length": 8192} at load time, because that becomes the effective limit.
• Skiplist. After changing model[2].skiplist_words, call model[2].resolve_with_tokenizer(model.tokenizer). Token ids are cached.
• Loss memory. With CachedMultiVectorMultipleNegativesRankingLoss, mini_batch_size bounds memory without changing results. mini_batch_num_tokens packs chunks by token budget instead, which suits variable-length documents (16 docs × ~940 tokens ≈ 15,000).
• Column order. Column 1 is encoded as the query and the remaining columns as documents. Override with router_mapping.
• Distillation data. KD data uses (query, doc_1..doc_N, scores). resolve_ids handles id-based datasets.

The full training script is on the page. Its model card declares apache-2.0.

Limits

• Single domain. Results cover one domain. The author says mLateOn-medical is not claimed to be strongest elsewhere.
• MIRIAD favours lexical matching. Its questions are generated from their own passages, so lexical overlap is unusually high:
• BM25 beat every sparse model and all but three dense models. Do not expect that to transfer.
• Against only the 10k gold passages, nearly every model scored above 0.97. Add distractors until your own eval spreads out.
• Score scales differ. The starting-point (50k corpus) and final (200k corpus) scores are not comparable.
• Long passages are the worst case. Index sizes reflect ~900-token passages. Natural Questions–style passages (~125 vectors each) start about 7x smaller.
• Pruning is a floor. The pruning used was naive, so the bottom rows of the quantization table are a floor, not the frontier.
• No pooling-aware training. Pooling was measured without pooling-aware training.