Sentence Transformers 6.0 MultiVectorEncoder: Training and Finetuning Guide

TL;DR

Hugging Face added a MultiVectorEncoder model type to Sentence Transformers v6.0 and released a step‑by‑step training script that finetunes a ColBERT‑style late‑interaction retriever (e.g., multi-vector-encoder/mLateOn-medical) in 14.5 hours on a single RTX 3090, achieving 0.9139 NDCG@10—the best score among 50+ dense, sparse, lexical, and multi‑vector baselines on a medical retrieval benchmark.

What are Multi‑Vector (Late‑Interaction) Models?

Multi‑vector models keep one embedding per token instead of compressing an entire text into a single vector.

  • Retrieval uses the MaxSim operator: each query token finds its best‑matching document token, and the scores are summed.
  • Token‑level matching preserves fine‑grained relevance signals that single‑vector models average away, typically yielding stronger retrieval at the cost of a larger index.
  • The companion post "Multi‑Vector (Late Interaction) Embedding Models with Sentence Transformers" details the architecture, encoding, scoring, and indexing.

"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 … keeps one small vector per token and scores a query against a document with the MaxSim operator." – Blog excerpt

Why Finetune a Multi‑Vector Model?

Finetuning adapts the model to the vocabulary, query style, and relevance notion of a specific domain (e.g., medical, legal, code). Key observations:

  • Domain signals are captured token‑by‑token, so modest in‑domain data yields large gains.
  • Most released checkpoints truncate documents at 180–512 tokens; on long passages (average 941 tokens in the MIRIAD medical set) truncation costs up to 0.24 NDCG@10.
  • Starting from an unsupervised checkpoint (pre‑contrastive but not yet supervised) consistently outperforms fully supervised checkpoints for domain adaptation.
  • The entire finetuning pipeline runs on a single consumer GPU, making domain‑specific retrievers accessible to most teams.

Training Components Overview

Component Role
Model MultiVectorEncoder instance – either a pre‑trained checkpoint or a fresh model built on a base transformer
Dataset datasets.Dataset or DatasetDict containing query‑document pairs (or other formats required by the loss)
Loss Contrastive in‑batch negatives (CachedMultiVectorMultipleNegativesRankingLoss) or knowledge‑distillation loss
Training Arguments MultiVectorEncoderTrainingArguments – batch size, learning rate, prompts, etc.
Evaluator MultiVectorInformationRetrievalEvaluator for NDCG@10, acc@1, etc.
Trainer MultiVectorEncoderTrainer orchestrates the above

Each section below can be read independently; the code snippets are complete and runnable with pip install -U "sentence-transformers[train]".

Model Selection and Configuration

Finetuning an Existing Checkpoint

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  # allow full‑length docs
)
# Remove built‑in length caps (e.g., 180‑512 tokens)
model[0].query_length = None
model[0].document_length = None
# Optional: skip punctuation tokens to shrink the index
import string
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)
  • The checkpoint already contains query/document marker tokens, projection head, and scoring skiplist.
  • Lifting caps lets the model ingest the full 1,400‑token medical passages.

Building a Fresh Model from a Base Transformer

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    model_kwargs={"torch_dtype": "float32"},
)
  • The pipeline adds a random token‑level projection (128‑dim) and the usual ColBERT modules.
  • Even strong dense backbones (e.g., Alibaba-NLP/gte-modernbert-base) reach near‑checkpoint performance after training on 25 k pairs.

Which Starting Point to Pick?

Empirical comparison on 25 k MIRIAD pairs:

Checkpoint Zero‑shot NDCG@10 After finetuning Δ
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 –
Takeaway: Unsup‑trained checkpoints adapt best; fully supervised checkpoints often regress.

Dataset Preparation

The trainer accepts any datasets.Dataset (Hub or local). Required format depends on the loss:

  • Label column (label or score) if the loss expects a supervision signal.
  • Input columns are ordered; the first column is treated as the query, subsequent columns as documents (unless overridden via router_mapping).

Loading from the Hub (MIRIAD example)

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: 4_467_542})
  • Each row provides a (question, passage_text) pair; passages average 941 tokens.

Local CSV/JSON Example

from datasets import load_dataset
dataset = load_dataset("csv", data_files="my_file.csv")
# or JSON
# dataset = load_dataset("json", data_files="my_file.json")
  • For custom preprocessing, use Dataset.from_dict.

Loss Function Choice

For question‑answer pairs, the recommended loss is in‑batch negatives with GradCache:

from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss
loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,  # chunk size for memory; effective batch = 128 in the blog
)
  • mini_batch_size controls memory; the effective contrastive batch size remains 128.
  • Scale parameter: keep the default scale=1.0 for MaxSim scores; using the dense default (scale=20.0) would saturate gradients.
  • For distillation, see MultiVectorDistillKLDivLoss (not used in the blog example).

Training Arguments

Key settings that yielded the best results:

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,   # effective batch via GradCache
    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",
)
  • Prompts must be supplied explicitly; the model does not apply stored markers automatically.
  • Leaving max_length unset ensures training sees the full document length.
  • A higher learning rate (1e-4) performed best after sweeping 5e‑6 → 2e‑4.

Evaluator for Retrieval

The most informative evaluator is MultiVectorInformationRetrievalEvaluator built from a held‑out query set and a corpus with distractors:

from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator
# Build corpus, queries, and relevance mapping (see blog for full loop)
evaluator = MultiVectorInformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="miriad-dev",
    batch_size=16,
)
  • Include hard distractor passages to avoid saturation; the blog added ~190 k random training passages to a 10 k gold set.

Full Training Script

The following script reproduces the mLateOn‑medical model:

import logging, string, 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.losses import CachedMultiVectorMultipleNegativesRankingLoss
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator

logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)

def main():
    # 1️⃣ Load unsupervised checkpoint and lift caps
    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)

    # 2️⃣ Load 1 M medical QA pairs
    train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train").select(range(1_000_000))

    # 3️⃣ Define loss (GradCache in‑batch negatives)
    loss = CachedMultiVectorMultipleNegativesRankingLoss(model=model, mini_batch_size=16)

    # 4️⃣ Light dev evaluator (500 queries)
    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:
            pid = f"p{len(passage_to_id)}"
            passage_to_id[row["passage_text"]] = pid
            corpus[pid] = row["passage_text"]
        if idx < 500:
            qid = f"q{idx}"
            queries[qid] = row["question"]
            relevant_docs[qid] = {passage_to_id[row["passage_text"]]}
    dev_evaluator = MultiVectorInformationRetrievalEvaluator(
        queries=queries, corpus=corpus, relevant_docs=relevant_docs, name="miriad-dev", batch_size=16
    )

    # 5️⃣ Training arguments (see previous section)
    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",
    )

    # 6️⃣ Trainer & training
    trainer = MultiVectorEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        loss=loss,
        evaluator=dev_evaluator,
    )
    trainer.train()

    # 7️⃣ Save and optionally push to Hub
    model.save_pretrained("models/mLateOn-medical/final")
    try:
        model.push_to_hub("mLateOn-medical")
    except Exception:
        logging.error("Failed to upload model: \n" + traceback.format_exc())

if __name__ == "__main__":
    main()
  • Runtime: 14.5 h on a single RTX 3090 (peak 17.5 GB VRAM).
  • Data efficiency: 100 k pairs (≈75 min) achieve within 0.012 NDCG@10 of the full‑million‑pair run.

Index Size and Optimization

Multi‑vector indexes are larger because each token yields a vector (≈878 vectors per 941‑token passage). Raw fp16 storage for 200 k passages ≈ 45 GB.

Token Pooling

HierarchicalTokenPooling(pool_factor=4) reduces vectors by a factor of 4 with < 0.0033 NDCG@10 loss.

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
embeddings = model.encode_document(passages, token_pooling=pooling)
  • ¼ of the vectors → ~11 GB, NDCG@10 ≈ 0.8991.

Quantization & Pruning (PLAID)

Using 1‑bit residual quantization and modest pruning:

Config Vectors kept Index size 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
  • Quantization yields > 13× reduction with < 0.02 NDCG loss, making multi‑vector retrieval comparable in storage to dense models.

Evaluation Results

Finetuned multi-vector-encoder/mLateOn-medical vs. 50+ baselines on a 200 k‑passage medical benchmark (1 k held‑out queries):

Model Family NDCG@10 acc@1
multi-vector-encoder/mLateOn-medical Multi‑vector (finetuned) 0.9139 0.849
lightonai/mLateOn Multi‑vector (zero‑shot) 0.8520 0.758
lightonai/GTE-ModernColBERT-v1 (cap lifted) Multi‑vector (zero‑shot) 0.8502 0.763
Qwen/Qwen3-Embedding-4B Dense (zero‑shot) 0.7817 0.669
voyageai/voyage-4-nano Dense (zero‑shot) 0.7563 0.638
BM25 Lexical 0.7501 0.641
naver/splade-v3 Sparse (zero‑shot) 0.6853 0.574
Takeaway: Late‑interaction models dominate when document length is long; finetuning adds a further +0.062 NDCG@10 over the strongest zero‑shot model.

Practical Takeaways

  1. Start from an unsupervised checkpoint (e.g., lightonai/mLateOn-unsupervised) for domain finetuning.
  2. Lift document length caps to match your data; otherwise you lose up to 0.24 NDCG@10 on long passages.
  3. Use GradCache (CachedMultiVectorMultipleNegativesRankingLoss) to achieve large effective batch sizes on a single GPU.
  4. Apply a punctuation skiplist to shrink the index by ~10 % with no quality loss.
  5. Invest in index compression (token pooling, PLAID quantization) to bring multi‑vector storage into the same ballpark as dense models.
  6. Even modest data (100 k pairs) yields strong performance, making the approach feasible for many organizations.

Additional Resources

  • Training examples: MIRIAD medical finetuning, MS MARCO contrastive & distillation, multimodal ColPali, PEFT LoRA adapters.
  • Documentation: Installation, quickstart, usage, custom model creation, pretrained model list, training & loss overviews, API reference.
  • Companion post (usage): Multi‑Vector (Late Interaction) Embedding Models with Sentence Transformers – covers encoding, indexing, and serving.

"The objection that multi‑vector indexes are too big does not survive a properly configured index." – Acknowledgements to Omar Khattab for quantization measurements.

Sources

Related