Training Overview

Tip

Using an AI coding agent (Claude Code, Codex, Cursor, Gemini CLI, …)?

hf skills add train-sentence-transformers [--claude] [--global]

And ask your agent to fine-tune a ColBERT-style multi-vector retrieval model for whatever task you have in mind.

Why Finetune?

Finetuning multi-vector (a.k.a. late-interaction or ColBERT-style) models heavily improves their retrieval performance on your specific domain: the vocabulary, the query style, and the notion of relevance all differ between e.g. web search, legal discovery, code search, and scientific literature review. Because queries and documents are matched token by token with MaxSim, multi-vector models pick up fine-grained domain signals that single-vector models tend to average away, and they typically respond very well to even modest amounts of in-domain finetuning data.

Also see Training Examples for training scripts for common real-world recipes that you can adopt.

Training Components

Training MultiVectorEncoder models involves between 4 to 6 components:

Model
Learn how to initialize the model for training.
Dataset
Learn how to prepare the data for training.
Loss Function
Learn how to prepare and choose a loss function.
Training Arguments
Learn which training arguments are useful.
Evaluator
Learn how to evaluate during and after training.
Trainer
Learn how to start the training process.

Model

Multi-vector models consist of a sequence of Modules, Multi-Vector Encoder specific Modules or Custom Modules, allowing for a lot of flexibility. If you want to further finetune an existing multi-vector model (e.g. it has a modules.json file), then you don’t have to worry about which modules are used:

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("lightonai/LateOn")

But if instead you want to train from another checkpoint, or from scratch, then these are the most common architectures you can use:

When you train from a base transformer model, the classic ColBERT architecture is the default: a Transformer producing contextualized token embeddings, a token-level Dense projecting each token down to the multi-vector dimension (classically 128), a MultiVectorMask computing the per-token scoring mask, and a token-level Normalize.

Documentation

from sentence_transformers import MultiVectorEncoder

# Loading in fp32 is preferred for training if your memory can handle it
model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})
# MultiVectorEncoder(
#   (0): Transformer({'transformer_task': 'feature-extraction', 'modality_config': {'text': {'method': 'forward', 'method_output_name': 'last_hidden_state'}}, 'module_output_name': 'token_embeddings', 'architecture': 'ModernBertModel'})
#   (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'module_input_name': 'token_embeddings', 'module_output_name': 'token_embeddings'})
#   (2): MultiVectorMask({'skiplist_words': [], 'skiplist_tasks': ['document'], 'keep_only_token_ids': None})
#   (3): Normalize({'module_input_name': 'token_embeddings', 'module_output_name': 'token_embeddings'})
# )

The fresh projection is randomly initialized, so training is required before this model is useful.

The classic ColBERT tokenization tricks are all left off by default: no [MASK] query expansion, no [Q] / [D] prefix tokens, no document length cap, and no punctuation skiplist. To reproduce the full classic ColBERT recipe, configure them explicitly:

from torch import nn

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.base.modules import Dense, Transformer
from sentence_transformers.multi_vector_encoder.modules import MultiVectorMask
from sentence_transformers.base.modules import Normalize
import string

transformer = Transformer(
    "answerdotai/ModernBERT-base",
    query_expansion={"strategy": "fixed", "length": 32},  # pad queries to 32 tokens with [MASK], truncate longer ones
    document_length=300,  # also truncate (not pad) documents to 300 tokens
    model_kwargs={"torch_dtype": "float32"},
)
dense = Dense(
    in_features=transformer.get_embedding_dimension(),
    out_features=128,
    bias=False,
    activation_function=nn.Identity(),
    module_input_name="token_embeddings",
)
mask = MultiVectorMask(skiplist_words=list(string.punctuation))  # exclude punctuation from document scoring
normalize = Normalize(module_input_name="token_embeddings")

model = MultiVectorEncoder(
    modules=[transformer, dense, mask, normalize],
    prompts={"query": "[Q] ", "document": "[D] "},
)

See Creating Custom Models for more details on the module pipeline, including how to add extra per-token feature channels.

Tip

Multimodal models require additional dependencies. Install them with e.g. pip install -U "sentence-transformers[image]" for image support. See Installation for all options.

ColPali-style models match text queries against page images, so retrieval skips OCR, layout parsing, and chunking entirely. Building one works exactly like building a text model: point MultiVectorEncoder at a multimodal embedding backbone and a freshly initialized token-level projection is appended, which is the construction that turned VLMs into ColPali and ColQwen in the first place.

Documentation

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder(
    "Qwen/Qwen3-VL-Embedding-2B",
    model_kwargs={"torch_dtype": "bfloat16"},
    processor_kwargs={"min_pixels": 28 * 28, "max_pixels": 600 * 600},
)

print([type(module).__name__ for module in model])
# ['Transformer', 'Dense', 'MultiVectorMask', 'Normalize']
print(model.modalities)
# ['text', 'image', 'video', 'message']

This is the same four-module stack a text backbone produces, so every multi-vector knob behaves identically. The projection is randomly initialized, so training is required before the model is useful. Image documents (PIL images, file paths, or URLs) go through the same encode_document path as text documents do, and the processor_kwargs above cap the patch budget per page, which is the main lever on both memory and index size.

An omnimodal backbone takes the identical route and additionally covers audio:

model = MultiVectorEncoder(
    "LCO-Embedding/LCO-Embedding-Omni-3B-2605",
    model_kwargs={"torch_dtype": "bfloat16"},
    trust_remote_code=True,
)

print(model.modalities)
# ['text', 'image', 'audio', 'video', 'message']

Already published transformers-native *ForRetrieval checkpoints (ColPali, ColQwen2, ColModernVBert) also load, detected from the architecture in their config.json, as a Transformer with transformer_task="retrieval" followed by a MultiVectorMask. Their projection and normalization live inside the transformers model, so no Dense and no Normalize are added, and their processor bakes in the query prefix and the visual prompt. That route exists to keep the existing checkpoints working, so prefer a multimodal backbone for anything new. See Transformers-native retrievers for those loading details.

See Training Examples > Multimodal for scripts covering both routes.

Dataset

The MultiVectorEncoderTrainer trains and evaluates using datasets.Dataset (one dataset) or datasets.DatasetDict instances (multiple datasets, see also Multi-dataset training).

If you want to load data from the Hugging Face Datasets, then you should use datasets.load_dataset():

Documentation

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/msmarco-bm25", "triplet", split="train")

print(train_dataset)
"""
Dataset({
    features: ['query', 'positive', 'negative'],
    num_rows: 502931
})
"""

Note

Many Hugging Face datasets that work out of the box with Sentence Transformers have been tagged with sentence-transformers, allowing you to easily find them by browsing to https://huggingface.co/datasets?other=sentence-transformers. We strongly recommend that you browse these datasets to find training datasets that might be useful for your tasks.

If you have local data in common file-formats, then you can load this data easily using datasets.load_dataset():

Documentation

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")

or:

from datasets import load_dataset

dataset = load_dataset("json", data_files="my_file.json")

If you have local data that requires some extra pre-processing, my recommendation is to initialize your dataset using datasets.Dataset.from_dict() and a dictionary of lists, like so:

Documentation

from datasets import Dataset

queries = []
positives = []
# Open a file, do preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "query": queries,
    "positive": positives,
})

Each key from the dictionary will become a column in the resulting dataset.

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:

  1. 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.

  2. 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.

Be sure to re-order your dataset columns with Dataset.select_columns if your columns are not ordered correctly. For example, if your dataset has ["good_answer", "bad_answer", "question"] as columns, then this dataset can technically be used with a loss that requires (anchor, positive, negative) triplets, but the good_answer column will be taken as the query, bad_answer as the positive document, and question as the negative document.

Additionally, if your dataset has extraneous columns (e.g. sample_id, metadata, source, type), you should remove these with Dataset.remove_columns as they will be used as inputs otherwise. You can also use Dataset.select_columns to keep only the desired columns.

There are two multi-vector specific conventions on top of this:

  • 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. This is the same multi-column convention as (query, positive, negative_1, ...), read positionally by MultiVectorDistillKLDivLoss. For KD datasets that store query / document IDs alongside separate text datasets (e.g. lightonai/ms-marco-en-bge), you can use resolve_ids() to resolve the IDs on the fly: it expands the stored ID list into the numbered document columns.

  • Positional query / document assignment: the first column is embedded as the query and all following columns as documents. This default can be overridden per column via the standard router_mapping training argument, mapping column names to "query" or "document".

Multimodal Datasets

Tip

Multimodal models require additional dependencies. Install them with e.g. pip install -U "sentence-transformers[image]" for image support. See Installation for all options.

MultiVectorEncoder datasets are not limited to text columns. With a ColPali-style retriever or a multimodal backbone (see the Model section), document columns can hold images as PIL images, file paths, or URLs. The rules from Dataset Format still apply, including the positional assignment: the first column is embedded as the query and every following column as a document, so image columns need no special handling.

A visual document retrieval dataset is therefore just a text query column plus one or more page image columns:

from datasets import load_dataset

train_dataset = load_dataset("tomaarsen/llamaindex-vdr-en-train-preprocessed", "train", split="train")
print(train_dataset)
"""
Dataset({
    features: ['query', 'image', 'negative_0', 'negative_1', 'negative_2', 'negative_3'],
    num_rows: 10000
})
"""

# Keep the query, the positive page and one hard negative page
train_dataset = train_dataset.select_columns(["query", "image", "negative_0"])

Query-page pairs work just as well: training/multimodal/finetuning_colqwen2.py finetunes a pretrained ColQwen2 on a two-column (query, image) dataset with MultiVectorMultipleNegativesRankingLoss, relying on the other pages in the batch as in-batch negatives.

The data collator automatically handles multimodal preprocessing via the model’s preprocess method, so no manual tokenization or image processing is needed. Note that MaxSim scores every query token against every image patch embedding, so a page image contributes many more vectors than a sentence of text does: CachedMultiVectorMultipleNegativesRankingLoss is useful here to keep the effective batch large while only a few samples are encoded at a time. See Training Examples > Multimodal for complete training scripts.

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. This is the core of the training process.

Sadly, there is no single loss function that works best for all use-cases. Instead, which loss function to use greatly depends on your available data and on your target task. See Dataset Format to learn what datasets are valid for which loss functions. Additionally, the Loss Overview will be your best friend to learn about the options.

Most loss functions can be initialized with just the MultiVectorEncoder that you’re training, alongside some optional parameters, e.g.:

from datasets import load_dataset
from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss

# Load a model to train/finetune
# Loading in fp32 is preferred for training if your memory can handle it
model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})

# Initialize the loss: in-batch negatives, scored with MaxSim
loss = MultiVectorMultipleNegativesRankingLoss(model=model)

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/msmarco-bm25", "triplet", split="train")
print(train_dataset)
"""
Dataset({
    features: ['query', 'positive', 'negative'],
    num_rows: 502931
})
"""

Training Arguments

The MultiVectorEncoderTrainingArguments class can be used to specify parameters for influencing training performance as well as defining the tracking/debugging parameters. Although it is optional, it is heavily recommended to experiment with the various useful arguments.

Key Training Arguments for improving training performance
learning_rate lr_scheduler_type warmup_steps num_train_epochs max_steps per_device_train_batch_size per_device_eval_batch_size auto_find_batch_size fp16 bf16 load_best_model_at_end metric_for_best_model gradient_accumulation_steps gradient_checkpointing eval_accumulation_steps optim batch_sampler multi_dataset_batch_sampler router_mapping learning_rate_mapping

Key Training Arguments for observing training performance
eval_strategy eval_steps save_strategy save_steps save_total_limit report_to run_name log_level logging_steps push_to_hub hub_model_id hub_strategy hub_private_repo

Here is an example of how MultiVectorEncoderTrainingArguments can be initialized:

args = MultiVectorEncoderTrainingArguments(
    # Required parameter:
    output_dir="models/multivector-ModernBERT-base-msmarco",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    learning_rate=3e-5,
    warmup_steps=0.05,
    fp16=False,  # Set to True if your GPU doesn't support BF16
    bf16=True,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # losses that use "in-batch negatives" benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.1,
    save_total_limit=2,
    logging_steps=0.01,
    run_name="multivector-ModernBERT-base-msmarco",  # Will be used in W&B if `wandb` is installed
)

Evaluator

You can provide the MultiVectorEncoderTrainer with an eval_dataset to get the evaluation loss during training, but it may be useful to get more concrete metrics during training, too. For this, you can use evaluators to assess the model’s performance with useful metrics before, during, or after training. You can use both an eval_dataset and an evaluator, one or the other, or neither. They evaluate based on the eval_strategy and eval_steps Training Arguments.

Here are the implemented Evaluators that come with Sentence Transformers for Multi-Vector Encoder models:

Evaluator

Required Data

MultiVectorInformationRetrievalEvaluator

Queries (qid => question), Corpus (cid => document), and relevant documents (qid => set[cid]).

MultiVectorNanoBEIREvaluator

No data required.

MultiVectorTripletEvaluator

(anchor, positive, negative) triplets.

MultiVectorRerankingEvaluator

List of {'query': '...', 'positive': [...], 'negative': [...]} dictionaries.

MultiVectorDistillationEvaluator

Queries with candidate documents and teacher scores.

Additionally, SequentialEvaluator should be used to combine multiple evaluators into one Evaluator that can be passed to the MultiVectorEncoderTrainer.

Sometimes you don’t have the required evaluation data to prepare one of these evaluators on your own, but you still want to track how well the model performs on some common benchmarks. In that case, you can use these evaluators with data from Hugging Face.

Documentation

from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator

# Initialize the evaluator. Unlike most other evaluators, this one loads the relevant datasets
# directly from Hugging Face, so there's no mandatory arguments
dev_evaluator = MultiVectorNanoBEIREvaluator()
# You can run evaluation like so:
# results = dev_evaluator(model)

Documentation

from datasets import load_dataset
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorTripletEvaluator

# Load triplets from the MS MARCO dataset (https://huggingface.co/datasets/sentence-transformers/msmarco-bm25)
max_samples = 1000
eval_dataset = load_dataset("sentence-transformers/msmarco-bm25", "triplet", split=f"train[:{max_samples}]")

# Initialize the evaluator
dev_evaluator = MultiVectorTripletEvaluator(
    anchors=eval_dataset["query"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    name="msmarco-dev",
)
# You can run evaluation like so:
# results = dev_evaluator(model)

Tip

When evaluating frequently during training with a small eval_steps, consider using a tiny eval_dataset to minimize evaluation overhead. If you’re concerned about the evaluation set size, a 90-1-9 train-eval-test split can provide a balance, reserving a reasonably sized test set for final evaluations. After training, you can assess your model’s performance using trainer.evaluate(test_dataset) for test loss or initialize a testing evaluator with test_evaluator(model) for detailed test metrics.

If you evaluate after training, but before saving the model, your automatically generated model card will still include the test results.

Warning

When using Distributed Training, the evaluator only runs on the first device, unlike the training and evaluation datasets, which are shared across all devices.

Trainer

The MultiVectorEncoderTrainer is where all previous components come together. We only have to specify the trainer with the model, training arguments (optional), training dataset, evaluation dataset (optional), loss function, evaluator (optional) and we can start training. Let’s have a look at a script where all of these components come together:

import logging

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 MultiVectorNanoBEIREvaluator
from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss

logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

# 1. Load a model to finetune with 2. (Optional) model card data
# Loading in fp32 is preferred for training if your memory can handle it
model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    model_card_data=MultiVectorEncoderModelCardData(
        language="en",
        license="apache-2.0",
        model_name="ColBERT ModernBERT-base trained on MS MARCO triplets",
    ),
    model_kwargs={"torch_dtype": "float32"},
)

# 3. Load a dataset to finetune on
full_dataset = load_dataset("sentence-transformers/msmarco-bm25", "triplet", split="train").select(range(51_000))
dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"]
eval_dataset = dataset_dict["test"]

# 4. Define a loss function: in-batch negatives, scored with MaxSim
loss = MultiVectorMultipleNegativesRankingLoss(model=model)

# 5. (Optional) Specify training arguments
run_name = "multivector-ModernBERT-base-msmarco"
args = MultiVectorEncoderTrainingArguments(
    # Required parameter:
    output_dir=f"models/{run_name}",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    learning_rate=3e-5,
    warmup_steps=0.05,
    fp16=False,  # Set to True if your GPU doesn't support BF16
    bf16=True,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
    load_best_model_at_end=True,
    metric_for_best_model="eval_NanoBEIR_mean_maxsim_ndcg@10",
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.1,
    save_total_limit=2,
    logging_steps=0.01,
    run_name=run_name,  # Will be used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = MultiVectorNanoBEIREvaluator(dataset_names=["msmarco", "nq", "fiqa2018"], batch_size=32)
dev_evaluator(model)

# 7. Create a trainer & train
trainer = MultiVectorEncoderTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# 8. Evaluate the model performance again after training
dev_evaluator(model)

# 9. Save the trained model
model.save_pretrained(f"models/{run_name}/final")

# 10. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(run_name)

The strongest late-interaction models are trained by distilling the rankings of a strong teacher (e.g. a Cross Encoder) over N candidate documents per query, rather than from raw pairs or triplets. lightonai/ms-marco-en-bge provides exactly that: per-query candidate document IDs with teacher scores, resolved to texts on the fly by resolve_ids().

See training_kd.py for a complete knowledge distillation run that additionally builds the model with a LateOn-style bottleneck projection head.

import logging

from datasets import load_dataset

from sentence_transformers import (
    MultiVectorEncoder,
    MultiVectorEncoderModelCardData,
    MultiVectorEncoderTrainer,
    MultiVectorEncoderTrainingArguments,
)
from sentence_transformers.util import resolve_ids
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator
from sentence_transformers.multi_vector_encoder.losses import MultiVectorDistillKLDivLoss

logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

max_list_length = 32  # Number of candidate documents (1 positive + negatives) scored per query
train_batch_size = 4  # Each batch holds train_batch_size * max_list_length documents, so keep it small

# 1. Load a model to finetune with 2. (Optional) model card data
# Loading in fp32 is preferred for training if your memory can handle it
model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    model_card_data=MultiVectorEncoderModelCardData(
        language="en",
        license="apache-2.0",
        model_name="ColBERT ModernBERT-base distilled from BGE on MS MARCO",
    ),
    model_kwargs={"torch_dtype": "float32"},
)

# 3. Load the KD dataset: per-query candidate ids + teacher scores, with separate text datasets
train_dataset = load_dataset("lightonai/ms-marco-en-bge", "train", split="train").select(range(20_000))
queries = load_dataset("lightonai/ms-marco-en-bge", "queries", split="train")
documents = load_dataset("lightonai/ms-marco-en-bge", "documents", split="train")

# resolve_ids resolves query_id -> query text and document_ids -> document texts on the fly.
train_dataset.set_transform(
    resolve_ids({"query_id": queries, "document_ids": documents}, max_list_length=max_list_length)
)

# 4. Define a loss function: KL divergence between the teacher and student score distributions
loss = MultiVectorDistillKLDivLoss(model=model)

# 5. (Optional) Specify training arguments
run_name = "multivector-ModernBERT-base-msmarco-kd"
args = MultiVectorEncoderTrainingArguments(
    # Required parameter:
    output_dir=f"models/{run_name}",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=train_batch_size,
    per_device_eval_batch_size=train_batch_size,
    learning_rate=3e-5,
    warmup_steps=0.05,
    fp16=False,  # Set to True if your GPU doesn't support BF16
    bf16=True,  # Set to True if you have a GPU that supports BF16
    load_best_model_at_end=True,
    metric_for_best_model="eval_NanoBEIR_mean_maxsim_ndcg@10",
    # Optional tracking/debugging parameters:
    eval_strategy="steps",  # The NanoBEIR evaluator runs on its own datasets, so no eval_dataset is needed
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.1,
    save_total_limit=2,
    logging_steps=0.01,
    run_name=run_name,  # Will be used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = MultiVectorNanoBEIREvaluator(dataset_names=["msmarco", "nq", "fiqa2018"], batch_size=train_batch_size)
dev_evaluator(model)

# 7. Create a trainer & train
trainer = MultiVectorEncoderTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# 8. Evaluate the model performance again after training
dev_evaluator(model)

# 9. Save the trained model
model.save_pretrained(f"models/{run_name}/final")

# 10. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(run_name)

Callbacks

This Multi-Vector Encoder trainer integrates support for various transformers.TrainerCallback subclasses, such as:

  • WandbCallback to automatically log training metrics to W&B if wandb is installed

  • TensorBoardCallback to log training metrics to TensorBoard if tensorboard is accessible.

  • CodeCarbonCallback to track the carbon emissions of your model during training if codecarbon is installed.

    • Note: These carbon emissions will be included in your automatically generated model card.

See the Transformers Callbacks documentation for more information on the integrated callbacks and how to write your own callbacks.

Multi-Dataset Training

The top performing models are trained using many datasets at once. Normally, this is rather tricky, as each dataset has a different format. However, MultiVectorEncoderTrainer can train with multiple datasets without having to convert each dataset to the same format. It can even apply different loss functions to each of the datasets. The steps to train with multiple datasets are:

  • Use a dictionary of Dataset instances (or a 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 function for different datasets.

Each training/evaluation batch will only contain samples from one of the datasets. The order in which batches are samples 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.

Training Tips

Multi-Vector Encoder models have a few quirks that you should be aware of when training them:

  1. The contrastive losses default to scale=1.0 (i.e. temperature=1.0), matching PyLate, and unlike the dense MultipleNegativesRankingLoss default of scale=20.0. That 20.0 exists to amplify bounded cosine similarity ([-1, 1]), but MaxSim is an unbounded sum over query-token similarities (range ~[0, num_query_tokens]), so it needs no amplification, exactly as the dense loss recommends scale=1 for dot-product similarity. A large scale here would saturate the softmax and kill gradients.

  2. The strongest late-interaction models are trained almost exclusively with n-way knowledge distillation from a stronger teacher model using MultiVectorDistillKLDivLoss, instead of training directly from text pairs or triplets. See the Knowledge Distillation tab under Trainer.

  3. In-batch negatives losses benefit heavily from larger batch sizes. If GPU memory is the bottleneck, CachedMultiVectorMultipleNegativesRankingLoss reaches much larger effective batch sizes at a small speed cost via GradCache.

  4. A fresh model from a base transformer starts without the classic ColBERT tokenization tricks ([MASK] query expansion, [Q] / [D] prefixes, a document length cap, a punctuation skiplist). They are worth configuring explicitly: query expansion and the prefix tokens in particular are part of the classic recipe that most released checkpoints use. See Creating Custom Models for the full set of defaults and how to change each one.

  5. Multi-vector models are evaluated (and scored during training) with MaxSim, and the per-query-token score contributions are inspectable: see the interpretability utilities for similarity maps and heatmaps on image documents.

Comparisons with SentenceTransformer Training

Training MultiVectorEncoder models is very similar to training SentenceTransformer models, with some key differences:

  • In SentenceTransformer training, a column is only encoded under a specific task if you say so with the router_mapping training argument. For MultiVectorEncoder training, the assignment is positional by default: the first column is encoded as the "query" and every following column as a "document", regardless of the column names. router_mapping still overrides this per column.

  • SentenceTransformer models produce one embedding per input, whereas MultiVectorEncoder models produce one embedding per token. Pairs are therefore scored with MaxSim rather than cosine similarity, which is why the contrastive losses default to scale=1.0 instead of scale=20.0. See Training Tips.

  • Multi-vector models carry tokenization knobs that have no dense equivalent: [MASK] query expansion, [Q] / [D] prefix tokens, a separate document length cap, and a punctuation skiplist. See Creating Custom Models.

See the Sentence Transformer > Training Overview documentation for more details on training SentenceTransformer models.