Sentence Transformers v6.0 Adds Multi-Vector Late Interaction Training and Compression

Sentence Transformers v6.0 Adds Multi-Vector Late Interaction Training and Compression Sentence Transformers v6.0 introduces MultiVectorEncoder for ColBERT-style late interaction retrieval and training. Using this framework, a specialized medical model trained on a single consumer GPU outperformed significantly larger dense models, zero-shot checkpoints, and BM25 on long-passage retrieval benchmarks. Sentence Transformers v6.0 enables ColBERT-style MultiVectorEncoder models that use token-level MaxSim scoring instead of single-vector dot products. A fine-tuned mLateOn-medical model trained on a single RTX 3090 achieved 0.9139 NDCG@10, outperforming dense baselines such as Qwen3-Embedding-4B and Qwen3-Embedding-8B. Unsupervised checkpoints adapted to domain fine-tuning significantly better than fully trained multi-vector checkpoints, which stagnated or regressed. Removing native document truncation caps improved multi-vector NDCG@10 by 0.08 to 0.24 on 941-token medical passages. Index compression methods like HierarchicalTokenPooling and fast-plaid reduced multi-vector storage from 45 GB down to 1.45 GB with minimal performance loss. Multi-Vector Architecture and Training Dynamics MultiVectorEncoder implements ColBERT-style late interaction retrieval, preserving one vector per token and computing query-document relevance through the MaxSim operator, which sums the best-matching document token scores for each query token. Training leverages CachedMultiVectorMultipleNegativesRankingLoss, using GradCache to decouple contrastive batch sizes from GPU memory limits by chunking document encodings. Multi-vector contrastive losses default to a scale factor of 1.0 rather than dense training's 20.0, because unscaled MaxSim scores naturally span roughly zero to the query length; adopting dense scaling values saturates the softmax function and eliminates gradient flow. Furthermore, contrastive batch sizes beyond 128 yielded no additional performance improvements in ablations. Ablations and Checkpoint Adaptation Behavior Evaluations revealed distinct architectural and adaptation behaviors during domain fine-tuning. Unsupervised pre-trained checkpoints (such as mLateOn-unsupervised) adapted substantially better to domain data than fully finished checkpoints, which plateaued or regressed across tested learning rates. A learning rate sweep from 5e-6 to 2e-4 identified 1e-4 as optimal. Setting a punctuation skiplist to exclude punctuation from document storage and scoring reduced index sizes by 9.6% while improving retrieval quality. Conversely, adding [MASK] query expansion produced no measurable performance impact, and restricting training sequence length to 512 tokens caused an NDCG@10 deficit of approximately 0.015 that persisted as training data increased. MIRIAD Training Implementation and Benchmark Results The multi-vector-encoder/mLateOn-medical model was trained on 1 million question-passage pairs from the MIRIAD dataset split using a single RTX 3090 GPU over 14.5 hours with a peak VRAM footprint of 17.5 GB. Scaling experiments indicated that 100,000 pairs trained in 75 minutes reached within 0.012 NDCG@10 of the full 1-million-pair run. When evaluated against 1,000 questions across a 200,000-passage corpus, mLateOn-medical reached 0.9139 NDCG@10 and 0.849 acc@1, surpassing its zero-shot base (0.8520 NDCG@10) by +0.062. It also outperformed Qwen3-Embedding-4B (0.7817 NDCG@10), despite the latter having roughly 33 times more active parameters, and exceeded BM25 (0.7501 NDCG@10). Lifting document length caps yielded gains up to +0.24 NDCG@10 across multi-vector models on the 941-token passages. Index Compression via Pooling and Fast-Plaid Storing one vector per token across the 200,000-passage corpus creates approximately 878 vectors per passage, totaling 45 GB in fp16 compared to under 1 GB for dense models. To mitigate this overhead, HierarchicalTokenPooling clusters document embeddings; keeping half the vectors decreased NDCG@10 by only 0.0033 with no loss in acc@1, while keeping one-quarter reduced index size to 11.2 GB at 0.8991 NDCG@10. Additionally, Omar Khattab's fast-plaid evaluation—using 1-bit residual quantization, 17-bit centroid IDs, and 18-bit document IDs—compressed the index to 3.37 GB at full vector retention (0.8984 NDCG@10) and 1.45 GB at 42% retention (0.8642 NDCG@10). At 1.45 GB, the index is smaller than the 1.64 GB uncompressed fp16 embeddings of Qwen3-Embedding-8B while scoring 0.0895 higher in NDCG@10.

Sentence Transformers v6.0 引入多向量嵌入模型训练与微调支持 Sentence Transformers v6.0 推出了支持 ColBERT 风格延迟交互检索与训练的 MultiVectorEncoder。实验表明，在单张 RTX 3090 显卡上微调得到的医疗领域多向量模型在专业检索评估中明显超越通用密集、稀疏、词法及预训练多向量模型，同时解除文档截断与使用索引压缩技术对性能与存储开销有关键影响。 Sentence Transformers v6.0 支持通过 pip install -U "sentence-transformers[train]" 运行 ColBERT 风格的 MultiVectorEncoder 训练与评估。 基于 lightonai/mLateOn-unsupervised 在 100 万医疗问答对上微调 14.5 小时得到的 mLateOn-medical，在 20 万文档库评估中达到 0.9139 NDCG@10，大幅领先通用密集模型和基线。 未监督预训练检查点比完全训练好的成品检查点更适应特定领域微调，后者在微调中常出现增益微弱或性能回退。 取消传统 ColBERT 和密集模型的长度截断上限可为长文本检索带来多达 +0.24 NDCG@10 的提升。 通过标点跳跃列表、分层 Token 池化以及 fast-plaid 量化技术，多向量索引体积可大幅压缩至与大体量密集模型相当甚至更小，同时保持较高检索质量。 延迟交互与多向量检索机制 传统密集模型将全文压缩为单个向量并通过点积打分，而多向量模型为每个 Token 保留一个向量，利用 MaxSim 算子累加每个查询 Token 与文档 Token 的最高相似度得分。这种 Token 级别的匹配提升了检索精度，但会显著增加索引体积。经典 ColBERT 和密集模型常将文档截断在 180 至 512 个 Token，但在平均 941 个 Token 的医疗段落评估中，截断导致了高达 0.24 NDCG@10 的损失。 领域微调策略与检查点选择 在 2.5 万对医疗问答评估中，未监督预训练模型适应领域微调的效果显著优于完全成品检查点。例如 lightonai/mLateOn-unsupervised 的 NDCG@10 提升了 0.0311，而 mLateOn 仅提升 0.0042，部分完成品模型甚至出现性能倒退。此外，利用随机初始化的投影层直接从基础模型（如 gte-modernbert-base）训练也能达到接近成熟检查点的表现，而 [MASK] 查询扩展对领域微调则未表现出可测量的性能差异。 训练优化与工程配置 训练采用了 CachedMultiVectorMultipleNegativesRankingLoss，利用 GradCache 将对比学习批次大小与显存解耦。实验显示有效对比批次超过 128 后无额外收益，且由于 MaxSim 得分是自然求和，损失函数的 scale 默认保持为 1.0（若沿用密集模型的 20.0 会导致梯度消失）。在单张 RTX 3090（峰值显存 17.5 GB）上，解除文档长度上限并配置学习率 1e-4、bf16，历时 14.5 小时即可完成 100 万样本训练；缩放实验表明 100k 样本训练 75 分钟的效果仅落后百万样本 0.012 NDCG@10。 检索评估与基准对比 在包含 1,000 个测试问题和 20 万段落（1 万目标段落加 19 万干扰段落）的 MIRIAD 评测中，mLateOn-medical 取得 0.9139 NDCG@10，超过未微调基线（0.8520）。相比同架构的密集模型，多向量模型具有明显优势；参数量为其约 33 倍的密集模型 Qwen3-Embedding-4B 仅取得 0.7817 NDCG@10。BM25 凭借无长度截断和测试集文本重合度获得 0.7501 NDCG@10，超过了部分截断的神经检索模型。 存储开销控制与量化压缩 针对 20 万段落保存全部多向量在 fp16 下需要约 45 GB。加入标点跳跃列表可在轻微提升效果的同时缩减 9.6% 索引体积。通过 HierarchicalTokenPooling 保留四分之一向量，可将索引降至 11.2 GB（0.8991 NDCG@10）。采用 fast-plaid（1 位残差量化）仅保留 42% 向量时，索引体积降至 1.45 GB（0.8642 NDCG@10），体积小于未压缩的 Qwen3-Embedding-8B（1.64 GB），而 NDCG@10 反高出 0.0895。

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. Its v6.0 update introduces a fourth model type: MultiVectorEncoder , for ColBERT-style late interaction retrieval, alongside a complete training approach for it. 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. Lastly, in the Evaluation section, I'll show you that my finetuned multi-vector-encoder/mLateOn-medical model, trained in 14.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. This blogpost is about training multi-vector models. If you want to learn how to use them, from loading and encoding to indexing in vector databases, see the companion Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers blogpost. Table of Contents What are Multi-Vector models? Why Finetune? Training Components Model Finetuning an existing multi-vector model Building one from a base transformer Which starting point should you pick? Dataset Data on the Hugging Face Hub Local Data Dataset Format Loss Function Training Arguments Evaluator Trainer Callbacks Multi-Dataset Training Evaluation Optimizing the index Acknowledgements Additional Resources Training Examples Documentation What are Multi-Vector models? 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. Why Finetune? 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. 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. On my medical evaluation with passages averaging 941 tokens, I measured that this truncation costs up to 0.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. Training Components Training MultiVectorEncoder models involves the following components: Model : The model to finetune or the architecture to build fresh. Dataset : The data used for training and evaluation. Loss Function : A function that measures the model's performance and guides the optimization process. Training Arguments (optional): Parameters that impact training performance, tracking, and debugging. Evaluator (optional): A class for evaluating the model before, during, or after training. Trainer : Brings together all training components. Let's take a closer look at each component. Model Multi-vector training gives you a real choice of starting point, and it matters more than you might expect. Finetuning an existing multi-vector model If you want to further finetune an existing multi-vector model, you don't have to worry about the architecture at all: from sentence_transformers import MultiVectorEncoder # Loading in fp32 is preferred for training if your memory can handle it model = MultiVectorEncoder( "lightonai/mLateOn-unsupervised" , model_kwargs={ "torch_dtype" : "float32" }, processor_kwargs={ "model_max_length" : 8192 }, # the tokenizer-level token limit ) 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. The first thing to check is the length configuration, since many released checkpoints cap documents at 180 to 512 tokens (see Why Finetune? ), 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: # Let the model read full documents instead of the caps it was trained with, # e.g. GTE-ModernColBERT-v1 ships with query_length=48 and document_length=300 model[ 0 ].query_length = None model[ 0 ].document_length = None 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. I made one more change, adding a punctuation skiplist that excludes punctuation tokens from document-side scoring and storage. In a 4-way ablation (none, punctuation, stopwords, both) it modestly won on quality, and it shrinks the document index by 9.6% on this data for free: import string # model[2] is the MultiVectorMask module model[ 2 ].skiplist_words = list (string.punctuation) model[ 2 ].resolve_with_tokenizer(model.tokenizer) # token ids are cached, so re-resolve after changing Building one from a base transformer You can also point MultiVectorEncoder at any base transformer, and a fresh, randomly initialized token-level projection is appended for you: from sentence_transformers import MultiVectorEncoder model = MultiVectorEncoder( "answerdotai/ModernBERT-base" , model_kwargs={ "torch_dtype" : "float32" }) # MultiVectorEncoder( # (0): Transformer({..., 'architecture': 'ModernBertModel'}) # (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, ...}) # (2): MultiVectorMask({'skiplist_words': [], 'skiplist_tasks': ['document'], ...}) # (3): Normalize({...}) # ) 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 . The projection starts random, so training is required before this model is useful. Interestingly, this works with strong dense embedding backbones too. A fresh projection on Alibaba-NLP/gte-modernbert-base reached within 0.03 of the existing-checkpoint starting points in my experiments, from nothing but the projection and 25k training pairs. The classic ColBERT tokenization tricks ( [MASK] query expansion, [Q] / [D] prefix tokens, a document length cap, a punctuation skiplist) are all off by default and configurable. See Creating Custom Models for the full set. For what it's worth, I tested [MASK] query expansion in four configurations for my domain finetune and none of them made a measurable difference, so don't feel obliged to reach for the classic recipe. Which starting point should you pick? I measured this directly while preparing this blogpost, taking six starting points and training each with the identical recipe on 25k medical question-passage pairs from MIRIAD , then evaluating on 1,000 held-out questions against a 50,000 passage corpus: Starting point Zero-shot NDCG@10 After 25k pairs Delta lightonai/mLateOn-unsupervised 0.9087 0.9398 +0.0311 lightonai/mLateOn 0.9277 0.9319 +0.0042 lightonai/LateOn-unsupervised 0.9026 0.9206 +0.0180 lightonai/LateOn 0.9185 0.9105 -0.0080 lightonai/GTE-ModernColBERT-v1 0.9198 0.9007 -0.0191 Fresh head on gte-modernbert-base - 0.9177 - The result surprised me, and it replicated across two model families. *The -unsupervised checkpoints adapt to a new domain far better than their finished siblings, overtaking them despite starting lower. These checkpoints sit after large-scale contrastive pretraining but before supervised finetuning on general retrieval, so they carry all the late-interaction structure with none of the general-purpose tuning that domain training then has to undo. The finished checkpoints, by contrast, barely moved or even regressed, at every learning rate I tried. So, if the model family you like publishes a pre-supervised checkpoint, start there. If not, a fresh projection on a strong retrieval-pretrained backbone is a close runner-up. Continuing from a fully finished checkpoint is the weakest option for domain adaptation, despite being the most natural-feeling one. Dataset The MultiVectorEncoderTrainer uses datasets.Dataset or datasets.DatasetDict instances for training and evaluation. You can load data from the Hugging Face Datasets Hub or use local data in whatever format you prefer (e.g. CSV, JSON, Parquet, Arrow, or SQL). Note: Lots of public datasets that work out of the box with Sentence Transformers have been tagged with sentence-transformers on the Hugging Face Hub, so you can easily find them on https://huggingface.co/datasets?other=sentence-transformers . Consider browsing through these to find ready-to-go datasets that might be useful for your tasks, domains, or languages. Data on the Hugging Face Hub You can use the load_dataset function to load data from datasets on the Hub: from datasets import load_dataset train_dataset = load_dataset( "tomaarsen/miriad-4.4M-split" , split= "train" ) print (train_dataset) """ Dataset({ features: ['question', 'passage_text'], num_rows: 4467542 }) """ This is the dataset I'll train on in this blogpost: 4.4 million medical questions from MIRIAD , each paired with the source passage that contains its answer (averaging 941 tokens). Simple (query, relevant passage) pairs like these are the easiest retrieval training data to collect for your own domain, and as you'll see, they're all you need. Local Data You can also use load_dataset for loading local data in common file formats: from datasets import load_dataset dataset = load_dataset( "csv" , data_files= "my_file.csv" ) # or dataset = load_dataset( "json" , data_files= "my_file.json" ) And if your local data requires pre-processing, you can use datasets.Dataset.from_dict to initialize your dataset with a dictionary of lists: from datasets import Dataset queries = [] documents = [] # Open a file, perform preprocessing, filtering, cleaning, etc. # and append to the lists dataset = Dataset.from_dict({ "query" : queries, "document" : documents, }) Dataset Format It is important that your dataset format matches your loss function (or that you choose a loss function that matches your dataset format). Verifying whether a dataset format works with a loss function involves two steps: If your loss function requires a Label according to the Loss Overview table, then your dataset must have a column named "label" or "score" . This column is automatically taken as the label. All columns not named "label" or "score" are considered Inputs according to the Loss Overview table. The number of remaining columns must match the number of valid inputs for your chosen loss. The names of these columns are irrelevant , only the order matters . There are two multi-vector specific conventions on top of this: Positional query and document assignment: the first column is embedded as the query and all following columns as documents , regardless of the column names. This default can be overridden per column via the standard router_mapping training argument. Knowledge distillation format: one column per candidate document, i.e. (query, document_1, ..., document_N, scores) where scores is a list of N teacher scores per row. For KD datasets that store query and document IDs alongside separate text datasets (e.g. lightonai/ms-marco-en-bge ), you can use resolve_ids to resolve the IDs to texts on the fly. Loss Function Loss functions quantify how well a model performs for a given batch of data, allowing an optimizer to update the model weights to produce more favourable (i.e., lower) loss values. The right loss function for your task depends on the data you have and what you're trying to achieve. You can find a full list of options in the Loss Overview . For the common case of question-answer or question-passage pairs, the workhorse is in-batch negatives training with MultiVectorMultipleNegativesRankingLoss , where every other document in the batch acts as a negative for each query. Bigger batches mean more negatives and stronger training, so in practice you'll want its GradCache variant, CachedMultiVectorMultipleNegativesRankingLoss , which decouples the effective batch size from what fits on your GPU: from sentence_transformers import MultiVectorEncoder from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss model = MultiVectorEncoder( "lightonai/mLateOn-unsupervised" , model_kwargs={ "torch_dtype" : "float32" }) loss = CachedMultiVectorMultipleNegativesRankingLoss( model=model, mini_batch_size= 16 , # how many documents to encode per chunk: bounds memory, not quality ) The mini_batch_size parameter bounds the memory by encoding documents in chunks of this size, while the effective contrastive batch size (128 in my run below, and in my ablations bigger batches bought nothing further) stays a free choice. GradCache guarantees identical results regardless of the chunk size, so lower it for smaller GPUs at only a wall-clock cost. When your document lengths vary a lot, consider its sibling mini_batch_num_tokens , which packs each chunk to a total token budget instead of a document count, so a chunk of unusually long documents can never spike your memory (my mini_batch_size=16 at roughly 940 tokens per document corresponds to mini_batch_num_tokens=15_000 ). One multi-vector specific trap is that the contrastive losses default to scale=1.0 , unlike the dense embedding equivalent which defaults to scale=20.0 . That 20.0 exists because a cosine similarity is a single value in [-1, 1], too narrow a range for a sharp softmax. A MaxSim score instead sums one best-match similarity per query token, so it already spans roughly [0, query_length]: a 32-token query can score up to 32. So don't copy scale=20.0 over from a dense training script, since it would saturate the softmax and kill your gradients. For distillation from a stronger teacher, which is how the strongest general-purpose late-interaction models are trained, see MultiVectorDistillKLDivLoss and the Knowledge Distillation tab in the Training Overview documentation. Training Arguments You can customize the training process using the MultiVectorEncoderTrainingArguments class. This class lets you adjust parameters that can impact training speed and help you understand what's happening during training. For more information on the most useful training arguments, check out the Multi-Vector Encoder > Training Overview > Training Arguments . It's worth reading to get the most out of your training. Here's an example, using the values from my actual training run: from sentence_transformers import MultiVectorEncoderTrainingArguments from sentence_transformers.base.sampler import BatchSamplers args = MultiVectorEncoderTrainingArguments( # Required parameter: output_dir= "models/mLateOn-medical" , # Optional training parameters: num_train_epochs= 1 , per_device_train_batch_size= 128 , # the effective contrastive batch, thanks to GradCache per_device_eval_batch_size= 16 , learning_rate= 1e-4 , warmup_steps= 0.05 , prompts={ "question" : "[Q] " , "passage_text" : "[D] " }, # the checkpoint's markers, keyed by training column fp16= False , # Set to True if you have a GPU that supports FP16 bf16= True , # Set to True if you have a GPU that supports BF16 batch_sampler=BatchSamplers.NO_DUPLICATES, # in-batch negatives benefit from no duplicates # Optional tracking/debugging parameters: eval_strategy= "steps" , eval_steps= 0.1 , save_strategy= "steps" , save_steps= 0.05 , logging_steps= 0.01 , run_name= "mLateOn-medical" , # Will be used in e.g. Trackio, W&B, etc. ) A few of these deserve a comment: prompts : training does not automatically apply the prompts stored in the model, so map them onto your training columns explicitly. Here that is the checkpoint's [Q] marker for the question column and [D] for the passage column, keeping training consistent with inference. max_length (deliberately not set): this argument caps tokenization during training only , for when you want cheaper training than the model's full serving length. I measured what that shortcut costs on this data. Training at 512 tokens lost about 0.015 NDCG@10 for about 2x the speed, and the deficit did not shrink with more data, because the model simply never sees what got cut off. Leave it unset so training matches inference, unless you need the speedup more than the quality. learning_rate=1e-4 : after a sweep from 5e-6 to 2e-4, I had the best luck with this higher-than-usual learning rate. Evaluator To track your model's performance during training, you can pass an eval_dataset to the trainer for evaluation loss, but concrete retrieval metrics are much more informative. Sentence Transformers includes the following built-in evaluators for multi-vector models: Evaluator Required Data MultiVectorInformationRetrievalEvaluator Queries, corpus, and relevant document mappings MultiVectorNanoBEIREvaluator No data required MultiVectorTripletEvaluator (anchor, positive, negative) triplets MultiVectorRerankingEvaluator List of {'query': '...', 'positive': [...], 'negative': [...]} dictionaries MultiVectorDistillationEvaluator Queries with candidate documents and teacher scores For domain finetuning, the MultiVectorInformationRetrievalEvaluator built from your own held-out data is the one that matters. One tip on constructing it is that the corpus should be hard enough that models can be told apart. In my case the MIRIAD questions are generated from their own source passages, which makes retrieval unusually easy. Against just the 10k gold passages, nearly every model scored above 0.97 NDCG@10. If your evaluation saturates like that, add distractor passages (I use deduplicated passages from the training split) until the scores spread out: from datasets import load_dataset from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator dataset = load_dataset( "tomaarsen/miriad-4.4M-split" ) # Gold: 1,000 evaluation questions, each mapping to its own passage, with the # eval split's full ~10k unique passages as the initial corpus corpus = {} queries = {} relevant_docs = {} passage_to_id = {} for idx, row in enumerate (dataset[ "eval" ]): if row[ "passage_text" ] not in passage_to_id: passage_to_id[row[ "passage_text" ]] = f"p { len (passage_to_id)} " corpus[passage_to_id[row[ "passage_text" ]]] = row[ "passage_text" ] if idx < 1_000 : queries[ f"q {idx} " ] = row[ "question" ] relevant_docs[ f"q {idx} " ] = {passage_to_id[row[ "passage_text" ]]} # Distractors: unique train passages that make the haystack realistic seen = set (passage_to_id) for row in dataset[ "train" ]: if len (corpus) >= 200_000 : break if row[ "passage_text" ] not in seen: seen.add(row[ "passage_text" ]) corpus[ f"d { len (corpus)} " ] = row[ "passage_text" ] evaluator = MultiVectorInformationRetrievalEvaluator( queries=queries, corpus=corpus, relevant_docs=relevant_docs, name= "miriad-dev" , batch_size= 16 , ) # results = evaluator(model) Trainer The MultiVectorEncoderTrainer is where all previous components come together. Here is the complete script that trained multi-vector-encoder/mLateOn-medical , the model from the introduction: import logging import string import traceback from datasets import load_dataset from sentence_transformers import ( MultiVectorEncoder, MultiVectorEncoderModelCardData, MultiVectorEncoderTrainer, MultiVectorEncoderTrainingArguments, ) from sentence_transformers.base.sampler import BatchSamplers from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss logging.basicConfig( format = "%(asctime)s - %(message)s" , datefmt= "%Y-%m-%d %H:%M:%S" , level=logging.INFO) def main (): # 1. Load the starting checkpoint: contrastively pretrained, not yet supervised # Loading in fp32 is preferred for training if your memory can handle it model = MultiVectorEncoder( "lightonai/mLateOn-unsupervised" , model_kwargs={ "torch_dtype" : "float32" }, processor_kwargs={ "model_max_length" : 8192 }, model_card_data=MultiVectorEncoderModelCardData( language= "en" , license= "apache-2.0" , model_name= "mLateOn finetuned on MIRIAD medical retrieval" , ), ) # 2. Lift the per-task length caps so training and inference see full medical passages model[ 0 ].query_length = None model[ 0 ].document_length = None # 3. Skip punctuation tokens during scoring: a small quality win and a 9.6% smaller index model[ 2 ].skiplist_words = list (string.punctuation) model[ 2 ].resolve_with_tokenizer(model.tokenizer) # 4. Load 1 million medical question-passage pairs train_dataset = load_dataset( "tomaarsen/miriad-4.4M-split" , split= "train" ).select( range ( 1_000_000 )) # 5. In-batch negatives with GradCache: large effective batch, memory-bounded chunks loss = CachedMultiVectorMultipleNegativesRankingLoss(model=model, mini_batch_size= 16 ) # 6. A light dev evaluator to watch progress during training: 500 held-out questions # against the eval split's ~10k unique passages. The full 200k protocol runs afterwards. eval_split = load_dataset( "tomaarsen/miriad-4.4M-split" , split= "eval" ) corpus, queries, relevant_docs, passage_to_id = {}, {}, {}, {} for idx, row in enumerate (eval_split): if row[ "passage_text" ] not in passage_to_id: passage_to_id[row[ "passage_text" ]] = f"p { len (passage_to_id)} " corpus[passage_to_id[row[ "passage_text" ]]] = row[ "passage_text" ] if idx < 500 : queries[ f"q {idx} " ] = row[ "question" ] relevant_docs[ f"q {idx} " ] = {passage_to_id[row[ "passage_text" ]]} dev_evaluator = MultiVectorInformationRetrievalEvaluator( queries=queries, corpus=corpus, relevant_docs=relevant_docs, name= "miriad-dev" , batch_size= 16 ) # 7. Training arguments, as discussed above run_name = "mLateOn-medical" args = MultiVectorEncoderTrainingArguments( output_dir= f"models/ {run_name} " , num_train_epochs= 1 , per_device_train_batch_size= 128 , per_device_eval_batch_size= 16 , learning_rate= 1e-4 , warmup_steps= 0.05 , prompts={ "question" : "[Q] " , "passage_text" : "[D] " }, fp16= False , # Set to True if you have a GPU that supports FP16 bf16= True , # Set to True if you have a GPU that supports BF16 batch_sampler=BatchSamplers.NO_DUPLICATES, eval_strategy= "steps" , eval_steps= 0.1 , save_strategy= "steps" , save_steps= 0.05 , logging_steps= 0.01 , run_name=run_name, ) # 8. Create a trainer & train trainer = MultiVectorEncoderTrainer( model=model, args=args, train_dataset=train_dataset, loss=loss, evaluator=dev_evaluator, ) trainer.train() # 9. Save the trained model model.save_pretrained( f"models/ {run_name} /final" ) # 10. (Optional) Push it to the Hugging Face Hub try : model.push_to_hub(run_name) except Exception: logging.error( f"Error uploading model to the Hugging Face Hub:\n {traceback.format_exc()} " ) if __name__ == "__main__" : main() That's the whole recipe: a pre-supervised checkpoint, a million domain pairs, in-batch negatives, full document length, and a higher-than-usual learning rate. The run took 14.5 hours on my single RTX 3090 at a peak of 17.5 GB VRAM, and every one of those choices was the winner of a measured comparison rather than a guess. For readers on smaller budgets, my scaling experiments put 100k pairs (75 minutes of training) within 0.012 NDCG@10 of the full million-pair run. Most of the gain comes in the first hour. Callbacks The MultiVectorEncoder trainer supports various transformers.TrainerCallback subclasses, including: WandbCallback for logging training metrics to W&B if wandb is installed TensorBoardCallback for logging training metrics to TensorBoard if tensorboard is accessible CodeCarbonCallback for tracking carbon emissions during training if codecarbon is installed Enable these via the report_to training argument, e.g. report_to=["wandb", "codecarbon"] , with the required dependencies installed. It defaults to "none" , and report_to="all" activates every integration whose dependency is installed. Refer to the Transformers Callbacks documentation for more information on these callbacks and how to create your own. Multi-Dataset Training Typically, top-performing general-purpose models are trained on multiple datasets simultaneously. However, this approach can be challenging due to the varying formats of each dataset. Fortunately, the MultiVectorEncoderTrainer allows you to train on multiple datasets without requiring a uniform format. Additionally, it provides the flexibility to apply different loss functions to each dataset. Here are the steps to train with multiple datasets at once: Use a dictionary of datasets.Dataset instances (or a datasets.DatasetDict ) as the train_dataset (and optionally also eval_dataset ). (Optional) Use a dictionary of loss functions mapping dataset names to losses. Only required if you wish to use different loss functions for different datasets. Each training/evaluation batch will only contain samples from one of the datasets. The order in which batches are sampled from the multiple datasets is defined by the MultiDatasetBatchSamplers enum, which can be passed to the MultiVectorEncoderTrainingArguments via multi_dataset_batch_sampler . Valid options are: MultiDatasetBatchSamplers.ROUND_ROBIN : Round-robin sampling from each dataset until one is exhausted. With this strategy, it's likely that not all samples from each dataset are used, but each dataset is sampled from equally. MultiDatasetBatchSamplers.PROPORTIONAL (default): Sample from each dataset in proportion to its size. With this strategy, all samples from each dataset are used and larger datasets are sampled from more frequently. Evaluation To find out where the finetuned model stands, I evaluated it against over 50 retrieval model configurations across four architecture families on the MIRIAD evaluation set, built exactly as in the Evaluator section above, with 1,000 held-out medical questions searching 200,000 unique passages (the 10k gold passages hidden among 190k deduplicated distractors from the training split). This corpus is four times the size of the 50,000-passage one from Which starting point should you pick? , so scores are not comparable between the two tables. The headline results, with the full table in the collapsible below: Model Family NDCG@10 multi-vector-encoder/mLateOn-medical (mine) Multi-vector, finetuned 0.9139 lightonai/mLateOn Multi-vector, zero-shot 0.8520 lightonai/GTE-ModernColBERT-v1 (cap lifted) Multi-vector, zero-shot 0.8502 Qwen/Qwen3-Embedding-4B Dense, zero-shot 0.7817 voyageai/voyage-4-nano Dense, zero-shot 0.7563 BM25 Lexical 0.7501 naver/splade-v3 Sparse, zero-shot 0.6853 The finetuned model tops the table, beating the strongest zero-shot model of any architecture by +0.062 NDCG@10. In other words, the strongest zero-shot model returns the right passage as the very first hit for 75.8% of the queries, while the finetuned model does so for 84.9%, cutting the rank-1 error by more than a third. The architecture pattern is just as clear, with the top of the table exclusively late interaction. On long documents, one vector per token beats one vector per document, even at matched training and matched backbones. DenseOn and LateOn share training data and architecture except for the head, and the late-interaction sibling wins by +0.12, with the multilingual pair (mDenseOn and mLateOn) replicating this at +0.13. Scale doesn't rescue single vectors either. Qwen3-Embedding-4B , the strongest dense model with roughly 33x the active (non-embedding) parameters of mine, still stops 0.13 short, and the 8B version scores lower than the 4B. BM25 also performs surprisingly well, beating every sparse model, every truncation-capped multi-vector model, and all but three dense models: the multi-billion Qwen3-Embedding-4B and 8B , and voyage-4-nano , which reads its full 32k token context to edge past by just 0.006. Don't expect that to transfer to your own data though. MIRIAD's questions are generated from the passages, so the lexical overlap between a query and its gold passage is far larger than in typical retrieval, and BM25's unlimited context length lets it use every one of those overlapping words while most neural checkpoints truncate. A BM25 baseline is cheap and always worth running, just don't count on this margin. The full field at a glance, sorted by score and colored by architecture family. Click to see the full evaluation table Model Family NDCG@10 acc@1 multi-vector-encoder/mLateOn-medical (mine) Multi-vector, finetuned 0.9139 0.849 lightonai/mLateOn Multi-vector 0.8520 0.758 lightonai/GTE-ModernColBERT-v1 @1024 Multi-vector 0.8502 0.763 lightonai/LateOn @1024 Multi-vector 0.8485 0.760 lightonai/mLateOn-unsupervised Multi-vector 0.8304 0.733 mixedbread-ai/mxbai-edge-colbert-v0-32m @1024 Multi-vector 0.8186 0.727 Qwen/Qwen3-Embedding-4B Dense 0.7817 0.669 Qwen/Qwen3-Embedding-8B Dense 0.7747 0.654 perplexity-ai/pplx-embed-v1-late-0.6b @1024 Multi-vector 0.7702 0.632 lightonai/ColBERT-Zero Multi-vector 0.7613 0.675 LiquidAI/LFM2.5-ColBERT-350M Multi-vector 0.7582 0.664 voyageai/voyage-4-nano Dense 0.7563 0.638 BM25 Lexical 0.7501 0.641 jinaai/jina-embeddings-v5-text-small-retrieval Dense 0.7470 0.620 Qwen/Qwen3-Embedding-0.6B Dense 0.7408 0.620 perplexity-ai/pplx-embed-v1-0.6b Dense 0.7384 0.615 mixedbread-ai/mxbai-edge-colbert-v0-32m Multi-vector 0.7350 0.639 mixedbread-ai/mxbai-edge-colbert-v0-17m Multi-vector 0.7271 0.631 answerdotai/answerai-colbert-small-v1 @512 Multi-vector 0.7264 0.615 lightonai/DenseOn @1024 Dense 0.7239 0.597 lightonai/mDenseOn @1024 Dense 0.7227 0.585 jinaai/jina-embeddings-v5-text-nano-retrieval Dense 0.7206 0.587 microsoft/harrier-oss-v1-0.6b Dense 0.7126 0.572 Alibaba-NLP/gte-modernbert-base Dense 0.7102 0.582 Snowflake/snowflake-arctic-embed-l-v2.0 Dense 0.7068 0.568 perplexity-ai/pplx-embed-v1-late-0.6b Multi-vector 0.7008 0.570 google/embeddinggemma-300m Dense 0.7000 0.563 lightonai/DenseOn Dense 0.6943 0.570 naver/splade-v3 Sparse 0.6853 0.574 ibm-granite/granite-embedding-small-english-r2 Dense 0.6813 0.546 naver/splade-v3-distilbert Sparse 0.6806 0.567 codefuse-ai/F2LLM-v2-0.6B Dense 0.6799 0.536 colbert-ir/colbertv2.0 @512 Multi-vector 0.6785 0.571 prithivida/Splade_PP_en_v1 Sparse 0.6755 0.577 lightonai/LateOn Multi-vector 0.6713 0.561 tomaarsen/embeddinggemma-300m-miriad-unsloth Dense, finetuned 0.6705 0.530 lightonai/LateOn-regularized Multi-vector 0.6673 0.554 lightonai/LateOn-unsupervised Multi-vector 0.6672 0.553 lightonai/GTE-ModernColBERT-v1 Multi-vector 0.6612 0.555 opensearch-project/opensearch-neural-sparse-encoding-v2-distill Sparse 0.6518 0.531 nomic-ai/nomic-embed-text-v1.5 (prompted) Dense 0.6387 0.498 mixedbread-ai/mxbai-embed-large-v1 Dense 0.6355 0.502 BAAI/bge-large-en-v1.5 Dense 0.6308 0.498 jinaai/jina-colbert-v2 @1024 Multi-vector 0.6218 0.504 nomic-ai/nomic-embed-text-v1.5 Dense 0.6203 0.487 answerdotai/answerai-colbert-small-v1 Multi-vector 0.6184 0.514 tomaarsen/splade-modernbert-base-miriad Sparse, finetuned 0.6142 0.473 NeuML/biomedbert-base-colbert Multi-vector 0.5963 0.463 BAAI/bge-base-en-v1.5 Dense 0.5930 0.454 BAAI/bge-small-en-v1.5 Dense 0.5881 0.457 sentence-transformers/all-mpnet-base-v2 Dense 0.5159 0.396 jinaai/jina-colbert-v2 Multi-vector 0.4992 0.401 mixedbread-ai/mxbai-colbert-large-v1 Multi-vector 0.4690 0.358 sentence-transformers/static-retrieval-mrl-en-v1 Dense 0.4614 0.323 sentence-transformers/all-MiniLM-L6-v2 Dense 0.4458 0.321 colbert-ir/colbertv2.0 Multi-vector 0.4347 0.346 Models marked @N are evaluated with their document length cap lifted to N tokens, since their native caps (180 to 512 tokens) would otherwise truncate the 941-token average passages. For every multi-vector model this lift was worth +0.08 to +0.24 NDCG@10 over the as-served row, and even the dense DenseOn gained +0.03 from the same treatment. Note that this does not mean that multi-vector-encoder/mLateOn-medical is the strongest model on all domains. It's simply the strongest in my domain. This is totally fine, as I just need this model to work well on my data. Don't underestimate the power of finetuning multi-vector models on your domain. Fourteen and a half hours on a single consumer GPU produced a model that no general-purpose retriever comes close to on this data, and the recipe is a single script with no teacher model and no mined negatives! Optimizing the index The fair objection to multi-vector retrieval is index size, and this domain is close to the worst case for it. Storing one vector per token, my model needs about 878 vectors per passage, so the 200,000-passage corpus takes roughly 45 GB at fp16, where a dense model needs well under 1 GB. Document length is what makes that gap so wide. The Natural Questions passages in the companion post average about 125 token vectors each, seven times fewer, so a corpus of short passages starts from a far smaller index than this one does. The HierarchicalTokenPooling module compresses exactly this by clustering each document's token embeddings and storing the cluster means, keeping roughly 1 / pool_factor of the vectors: from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling pooling = HierarchicalTokenPooling(pool_factor= 4 ) document_embeddings = model.encode_document(passages, token_pooling=pooling) I measured it post-hoc on the finished model, with no pooling-aware training, and on long documents it is remarkably cheap. The solid points are uncompressed embeddings, so that every family is counted the same way and scored with exact search. You would not deploy any of them like that, though. Dense indexes routinely use int8 or binary quantization with rescoring, sparse indexes compress their postings, and multi-vector indexes use PLAID-style residual compression. Don't read those points as the disk you need to buy, but as relative storage cost. Token pooling is the solid line. Halving the vector count costs 0.0033 NDCG@10 and leaves rank-1 accuracy untouched, and keeping only a quarter of them, at 11.2 GB, still scores 0.8991. The curve keeps going (I measured out to a tenth of the vectors, still at 0.8765) but there is little reason to push pooling that far once quantization is on the table, which is what the dashed line below is about. The dashed line is what a real deployment might look like. I gave Omar Khattab early access to the model and the benchmark, and he measured these configurations with fast-plaid at 1-bit residual quantization, using compact 17-bit centroid ids and 18-bit document ids instead of its ordinary unpacked 64-bit integers, plus document-side pruning: configuration vectors kept index NDCG@10 1-bit PLAID, all vectors 100% 3.37 GB 0.8984 1-bit PLAID + pruning 65% 2.23 GB 0.8830 1-bit PLAID + pruning 42% 1.45 GB 0.8642 That first row is 13x smaller than the raw embeddings, for 0.0155 NDCG@10. That is a far better trade than anywhere on the pooling curve. Quantization shrinks each vector while pooling and pruning cut how many you keep, so they compose, and quantization is the one to reach for first. Push further and the last row lands at 1.45 GB, smaller than the fp16 embeddings of Qwen3-Embedding-8B (1.64 GB), while scoring 0.0895 higher. The objection that multi-vector indexes are too big does not survive a properly configured index. The pruning here is naive, meant only to establish that token reduction works on top of quantization, so read the bottom two rows as a floor rather than the frontier. If you would rather not hand-tune quantization at all, the Indexing section of the companion post covers fast-plaid, Qdrant, Weaviate, and Vespa. Multi-vector retrieval is only as expensive as its index. The raw embeddings for this corpus are 45 GB, and a properly configured index is at least 7x smaller at nearly the same accuracy. The index deserves as much of your attention as the checkpoint. Acknowledgements Thanks to Omar Khattab for measuring the quantized and pruned index configurations in Optimizing the index , and for the discussions around late-interaction index costs. Additional Resources Training Examples These pages have training examples with explanations as well as links to training scripts. You can use them to get familiar with the multi-vector training loop: MIRIAD : domain-specific training on medical retrieval, an earlier and simpler cousin of this blogpost's recipe MS MARCO : contrastive and knowledge distillation recipes Multimodal : ColPali-style visual document retrieval training PEFT Adapters : parameter-efficient finetuning with LoRA Documentation For further learning, you may also want to explore the following resources on Sentence Transformers: Installation Quickstart Usage Creating Custom Models Pretrained Models Training Overview (This blogpost is a distillation of the Training Overview documentation) Loss Overview API Reference And here is an advanced page that might interest you: Distributed Training And the companion blogpost, covering everything about using these models: Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers