Training and Fine-Tuning Multi-Vector Embedding Models with Sentence Transformers

Training and Fine-Tuning Multi-Vector Embedding Models with Sentence Transformers

Published August 26, 2026

Sentence Transformers is a Python library for using and training embedding and reranker models for applications including retrieval-augmented generation, semantic search, and semantic textual similarity. Version 6.0 adds MultiVectorEncoder, a fourth model type designed for ColBERT-style late-interaction retrieval, together with tools for training and fine-tuning these models.

This guide explains the main components of multi-vector training: models, datasets, loss functions, training arguments, evaluators, and trainers. The examples use pip install -U "sentence-transformers[train]".

The example model, multi-vector-encoder/mLateOn-medical, was trained in 14.5 hours on a single RTX 3090. On the MIRIAD medical retrieval evaluation, it outperformed the general-purpose dense, sparse, lexical, and multi-vector retrieval models tested.

What Are Multi-Vector Models?

A dense embedding model represents an entire text with one vector. Comparing two texts then requires a single dot product between those vectors.

A multi-vector model, also known as a late-interaction or ColBERT-style model, retains one small vector for each token. During retrieval, the MaxSim operation finds the best matching document token for every query token and sums those scores. This token-level matching preserves details that a single document vector may average away, generally improving retrieval quality at the cost of a larger index.

Why Fine-Tune?

Fine-tuning adapts retrieval to a specific domain's vocabulary, query style, and definition of relevance. These characteristics differ among web search, legal discovery, code search, scientific literature, and internal company documents. Because multi-vector models compare tokens individually, even modest amounts of in-domain data can improve their handling of domain-specific signals.

Many existing retrieval checkpoints were configured for short passages. Classic ColBERT checkpoints may truncate documents at 180 or 300 tokens, while popular dense models often use limits of 256 or 512 tokens. On a medical evaluation with passages averaging 941 tokens, truncation reduced NDCG@10 by as much as 0.24. Training a model for a particular domain allows its document length to be configured for that data.

Training Components

A MultiVectorEncoder training workflow contains these components:

1. Model: An existing model to fine-tune or a new architecture built from a transformer.
2. Dataset: Training and evaluation data.
3. Loss function: Measures model performance and guides optimization.
4. Training arguments: Optional settings for performance, tracking, and debugging.
5. Evaluator: An optional component for measuring retrieval quality before, during, or after training.
6. Trainer: Combines the model, data, loss, arguments, and evaluator.

Model

There are two main starting points: an existing multi-vector checkpoint or a base transformer with a newly initialized projection layer.

Fine-Tuning an Existing Multi-Vector Model

from sentence_transformers import MultiVectorEncoder

## Loading in fp32 is preferred for training if memory allows it
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},
)

An existing checkpoint includes its query and document marker tokens, projection head, and scoring skiplist. These components can usually remain unchanged while length settings are adapted to the training data.

For example, the mLateOn family supports the backbone's full 8192-token context. If a checkpoint has task-specific limits, they can be removed:

## Remove task-specific document and query limits
model[0].query_length = None
model[0].document_length = None

With these limits unset, truncation falls back to the tokenizer's model_max_length.

A punctuation skiplist can also exclude punctuation tokens from document-side scoring and storage:

import string

## model[2] is the MultiVectorMask module
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)

In a four-way ablation, the punctuation skiplist produced a modest quality improvement and reduced the document index by 9.6% on the medical data.

Building a Model from a Base Transformer

MultiVectorEncoder can also use a base transformer and add a randomly initialized token-level projection:

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    model_kwargs={"torch_dtype": "float32"},
)

The resulting pipeline contains:

• A Transformer that produces contextualized token embeddings.
• A token-level Dense layer that projects embeddings to 128 dimensions.
• A MultiVectorMask module that determines which tokens participate in scoring.
• A token-level Normalize module.

Because the projection is randomly initialized, the model must be trained before use. In the reported experiments, a new projection on Alibaba-NLP/gte-modernbert-base came within 0.03 of existing-checkpoint starting points after training on 25,000 pairs.

Classic ColBERT options, including [MASK] query expansion, [Q] and [D] prefix tokens, document length limits, and punctuation skipping, are configurable and disabled by default. In the reported domain fine-tuning tests, [MASK] query expansion did not produce a measurable difference.

Choosing a Starting Point

Six starting points were trained with the same recipe on 25,000 medical question-passage pairs from MIRIAD. Evaluation used 1,000 held-out questions and a 50,000-passage corpus.

The unsupervised checkpoints adapted to the new domain more effectively than their fully supervised counterparts. They had undergone large-scale contrastive pretraining but not supervised fine-tuning for general retrieval, allowing domain-specific training to build on their late-interaction structure without first undoing as much general-purpose tuning.

When a model family provides a pre-supervised checkpoint, it is a useful starting point for domain adaptation. A new projection on a strong retrieval-pretrained backbone is another option. Continuing from a fully finished checkpoint produced the weakest adaptation results in these experiments.

Dataset

MultiVectorEncoderTrainer accepts datasets.Dataset and datasets.DatasetDict objects for training and evaluation. Data can come from the Hugging Face Datasets Hub or local CSV, JSON, Parquet, Arrow, or SQL sources.

Loading Data from the Hugging Face Hub

from datasets import load_dataset

train_dataset = load_dataset(
    "tomaarsen/miriad-4.4M-split",
    split="train",
)

print(train_dataset)

The MIRIAD dataset contains 4,467,542 medical questions paired with source passages. The passages average 941 tokens. Simple question-passage pairs are sufficient for the contrastive training example used here.

Loading Local Data

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
## or
dataset = load_dataset("json", data_files="my_file.json")

For custom preprocessing, a dataset can be created from dictionaries of lists:

from datasets import Dataset

queries = []
documents = []
## Read, preprocess, filter, or clean data here

dataset = Dataset.from_dict({
    "query": queries,
    "document": documents,
})

Dataset Format

The dataset format must match the selected loss function:

1. If the loss requires a label, the dataset must contain a column named label or score. That column is automatically treated as the label.
2. All other columns are treated as inputs. Their names do not matter, but their order and number must match the loss function's expected inputs.

Multi-vector training adds two conventions:

• The first input column is embedded as the query, while subsequent columns are treated as documents. This can be overridden with the router_mapping training argument.
• Knowledge-distillation data uses one query, multiple candidate documents, and teacher scores: (query, document_1, ..., document_N, scores). When data stores query and document IDs separately from their text, resolve_ids can resolve those IDs during processing.

Loss Function

Loss functions measure model performance for each batch and guide weight updates. The appropriate choice depends on the available data and training objective.

For question-answer and question-passage pairs, MultiVectorMultipleNegativesRankingLoss uses every other document in the batch as a negative for each query. Larger batches provide more negatives. The cached variant, CachedMultiVectorMultipleNegativesRankingLoss, allows the effective batch size to exceed the amount that fits in GPU memory:

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,
)

mini_batch_size controls the number of documents encoded at a time, limiting memory use without changing the effective contrastive batch size. In the reported run, the effective batch size was 128. When document lengths vary substantially, mini_batch_num_tokens can limit each chunk by total token count instead of document count. A mini_batch_size of 16 with documents averaging about 940 tokens corresponds approximately to mini_batch_num_tokens=15_000.

Multi-vector contrastive losses use a default scale=1.0, unlike the dense embedding equivalent, which defaults to scale=20.0. Dense cosine similarity occupies a narrow range from -1 to 1, while MaxSim sums one similarity per query token and can reach approximately the query length. For this reason, copying scale=20.0 from a dense training script can saturate the softmax and eliminate useful gradients.

For teacher-based training, MultiVectorDistillKLDivLoss supports knowledge distillation from a stronger model.

Training Arguments

MultiVectorEncoderTrainingArguments controls training speed, evaluation, saving, logging, and debugging:

from sentence_transformers import MultiVectorEncoderTrainingArguments
from sentence_transformers.base.sampler import BatchSamplers

args = MultiVectorEncoderTrainingArguments(
    output_dir="models/mLateOn-medical",
    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,
    bf16=True,
    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="mLateOn-medical",
)

Important settings include:

• prompts: Model prompts are not automatically applied during training. The question and passage columns must be mapped explicitly to [Q] and [D] so training matches inference.
• max_length: This setting limits tokenization during training only. Training at 512 tokens was about twice as fast but reduced NDCG@10 by approximately 0.015 because the model never saw the truncated content.
• learning_rate=1e-4: A sweep from 5e-6 to 2e-4 produced the best results with this relatively high learning rate.

Evaluator

Retrieval metrics provide more useful feedback than training loss alone. Sentence Transformers includes these multi-vector evaluators:

For domain fine-tuning, MultiVectorInformationRetrievalEvaluator built from held-out domain data is the most relevant choice.

The evaluation corpus should be difficult enough to distinguish models. MIRIAD questions are generated from their source passages, making retrieval against only the approximately 10,000 gold passages unusually easy. Most models scored above 0.97 NDCG@10 in that setting. Adding deduplicated passages from the training split created a more realistic distractor corpus.

from datasets import load_dataset
from sentence_transformers.multi_vector_encoder.evaluation import (
    MultiVectorInformationRetrievalEvaluator,
)

dataset = load_dataset("tomaarsen/miriad-4.4M-split")

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"]]
        }

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,
)

Trainer

MultiVectorEncoderTrainer combines the model, dataset, loss, arguments, and evaluator. The following example represents the complete training setup for multi-vector-encoder/mLateOn-medical:

import string

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,
)

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",
    ),
)

model[0].query_length = None
model[0].document_length = None
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)

train_dataset = load_dataset(
    "tomaarsen/miriad-4.4M-split",
    split="train",
).select(range(1_000_000))

loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,
)

## Build a development evaluator from held-out data
## The full evaluation protocol uses 1,000 questions and 200,000 passages.
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,
)

args = MultiVectorEncoderTrainingArguments(
    output_dir="models/mLateOn-medical",
    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,
    bf16=True,
    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="mLateOn-medical",
)

trainer = MultiVectorEncoderTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)

trainer.train()
model.save_pretrained("models/mLateOn-medical/final")

The complete run used one million domain pairs, in-batch negatives, full document length, and a learning rate of 1e-4. It took 14.5 hours on one RTX 3090 and reached a peak of 17.5 GB of VRAM. A smaller run with 100,000 pairs took 75 minutes and came within 0.012 NDCG@10 of the full run.

Callbacks

The trainer supports standard transformers.TrainerCallback integrations, including:

• WandbCallback for Weights & Biases logging.
• TensorBoardCallback for TensorBoard logging.
• CodeCarbonCallback for tracking carbon emissions.

These can be enabled through the report_to argument, for example report_to=["wandb", "codecarbon"]. The default is "none", while report_to="all" enables installed integrations.

Multi-Dataset Training

MultiVectorEncoderTrainer can train on multiple datasets without requiring identical formats. A dictionary of datasets is supplied as train_dataset, and a dictionary of loss functions can optionally map dataset names to different losses.

Each batch contains examples from one dataset. Sampling is controlled by MultiDatasetBatchSamplers:

• ROUND_ROBIN: Samples datasets in turn until one is exhausted. Datasets are sampled equally, but not all examples may be used.
• PROPORTIONAL: The default. Samples in proportion to dataset size and uses all examples.

Evaluation Results

The fine-tuned model was evaluated on MIRIAD with 1,000 held-out medical questions searching 200,000 unique passages. The corpus contained approximately 10,000 gold passages and 190,000 deduplicated distractors from the training split.

The 200,000-passage corpus differs from the 50,000-passage corpus used for starting-point comparisons, so the scores are not directly comparable.

The fine-tuned model achieved the highest score, exceeding the strongest zero-shot model by 0.062 NDCG@10. Rank-1 accuracy was 0.849 for the fine-tuned model and 0.758 for the strongest zero-shot model.

The results also favored late interaction on these long documents. DenseOn and LateOn use shared training data and architecture apart from their output heads, and LateOn performed 0.12 better. The multilingual mDenseOn and mLateOn pair showed a similar 0.13 difference.

BM25 performed well in this evaluation because MIRIAD questions are generated from their source passages, producing substantial lexical overlap. BM25 also uses the full document, while many neural checkpoints truncate documents. This makes BM25 an important baseline, but the result may not transfer to domains with less query-document word overlap.

Models marked @N were evaluated with the document length cap lifted to N tokens. For the multi-vector models tested, lifting the cap improved NDCG@10 by 0.08 to 0.24 over the as-served configuration. DenseOn improved by 0.03 under the same treatment.

The results describe performance on the medical domain used for training and evaluation. They do not establish that multi-vector-encoder/mLateOn-medical is the strongest model for other domains.

Optimizing the Index

Multi-vector retrieval stores one vector per token, so index size depends heavily on document length. In this experiment, each passage required approximately 878 vectors. Storing the 200,000-passage corpus in fp16 required roughly 45 GB, while a dense index required well under 1 GB.

The average Natural Questions passage in the companion example contains approximately 125 token vectors, demonstrating why short passages produce considerably smaller multi-vector indexes.

HierarchicalTokenPooling reduces storage by clustering token embeddings and retaining cluster means:

from sentence_transformers.multi_vector_encoder.modules import (
    HierarchicalTokenPooling,
)

pooling = HierarchicalTokenPooling(pool_factor=4)
document_embeddings = model.encode_document(
    passages,
    token_pooling=pooling,
)

Post-hoc measurements on the finished model showed that pooling can reduce storage with limited quality loss. Halving the vector count reduced NDCG@10 by 0.0033 without changing rank-1 accuracy. Retaining one quarter of the vectors produced an 11.2 GB index with an NDCG@10 score of 0.8991. Retaining one tenth of the vectors produced a score of 0.8765.

Quantization and pruning reduced the index further. Using fast-plaid, 1-bit residual quantization, compact 17-bit centroid IDs, 18-bit document IDs, and document-side pruning produced these configurations:

The first configuration was 13 times smaller than the raw embeddings, with a 0.0155 NDCG@10 reduction. The final configuration used 1.45 GB, less than the 1.64 GB fp16 embeddings of Qwen3-Embedding-8B, while scoring 0.0895 higher in this evaluation.

The pruning experiment was intended to show how token reduction works alongside quantization rather than to identify an optimal deployment configuration. Quantization, pooling, and pruning can be combined, and the index configuration should be evaluated alongside retrieval quality.

Additional Resources

Sentence Transformers provides examples for:

• MIRIAD domain-specific medical retrieval training.
• MS MARCO contrastive and knowledge-distillation training.
• Multimodal ColPali-style visual document retrieval.
• PEFT and LoRA fine-tuning.

Related documentation covers installation, quickstarts, multi-vector usage, custom models, pretrained models, training, losses, the API, and distributed training. A companion guide covers using multi-vector models for encoding, indexing, and retrieval.