Sentence Transformers v6.0 Multi-Vector Encoder release

TL;DR

Sentence Transformers 6.0 introduces a fourth model type, MultiVectorEncoder, which implements ColBERT‑style late‑interaction (multi‑vector) retrieval for text and multimodal data, delivering stronger semantic search and visual document retrieval while requiring larger token‑level indexes.

What are Multi‑Vector Models?

A multi‑vector model keeps one embedding per token instead of compressing an entire passage into a single vector. After encoding, a query is scored against a document with the MaxSim operator, which sums the highest cosine similarity for each query token across all document tokens. This preserves exact token matches (e.g., product codes) and captures semantic paraphrases because token embeddings are contextualized.

"Encode ‘Where do penguins live?’ against ‘Penguins inhabit Antarctica.’ and the query token live finds its best match on inhabit with a similarity of 0.94." – Hugging Face blog

The MaxSim Operator

For query token set (Q) and document token set (D):

$$ \text{MaxSim}(Q, D) = \sum_{q_i \in Q} \max_{d_j \in D} ; q_i \cdot d_j $$

Because embeddings are L2‑normalized, each dot product is a cosine similarity in ([-1, 1]), so the total score lies in ([-|Q|, |Q|]).

Trade‑offs

  • Quality gain – token‑level matching improves retrieval for queries that rely on a single exact term or multiple independent constraints.
  • Index cost – one vector per token inflates storage (e.g., 4,874 passages require 311 MB raw for a 128‑dim model vs. 7.5 MB for a 384‑dim dense model). Compression techniques such as PLAID or hierarchical token pooling can reduce this overhead.

Installation

pip install -U sentence-transformers
# For visual document retrieval add image extras
pip install -U "sentence-transformers[image]"

Sentence Transformers v6.0 requires transformers v5.x, torch ≥ 2.2, and huggingface‑hub v1.x.

Loading a Multi‑Vector Model

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("lightonai/LateOn")

Any checkpoint tagged with multi-vector and sentence-transformers loads directly, regardless of whether it originated from PyLate, Stanford‑NLP ColBERT, or colpali‑engine. For ColPali visual models a small repository configuration is required (see Supported Models).

Inspecting Model Configuration

model = MultiVectorEncoder("colbert-ir/colbertv2.0")
print(model)
print(model.prompts)

Typical output shows a Transformer, a token‑level Dense projection (e.g., 128‑dim), a MultiVectorMask for skip‑list tokens, and a Normalize layer. Query and document length caps (e.g., 32 and 180 tokens) are also displayed.

Encoding Queries and Documents

Multi‑vector models are asymmetric; use the dedicated methods:

queries = ["What is the capital of France?"]
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
q_emb = model.encode_query(queries)          # shape (n_query_tokens, dim)
d_emb = model.encode_document(documents)    # list of (n_doc_tokens, dim)

Each document returns a matrix whose first dimension equals its token count, so tensors cannot be stacked into a single rectangular batch without padding.

Scoring with MaxSim

scores = model.similarity(q_emb, d_emb)  # all‑pairs MaxSim matrix
print(scores)

The result is a tensor of summed cosine similarities. Because the sum scales with query length, scores are not directly comparable across models with different query caps. To obtain a bounded metric, switch to MeanMaxSim:

model = MultiVectorEncoder("lightonai/LateOn", similarity_fn_name="meanmaxsim")
print(model.similarity(q_emb, d_emb))  # values in [0, 1]

Semantic Search (Exhaustive)

For small corpora you can encode the entire collection once and run exhaustive MaxSim at query time:

from datasets import load_dataset
from sentence_transformers import MultiVectorEncoder

corpus = list(dict.fromkeys(load_dataset("sentence-transformers/natural-questions", split="train[:5000]")["answer"]))
model = MultiVectorEncoder("lightonai/LateOn")
corpus_emb = model.encode_document(corpus, convert_to_tensor=True)

query = "when did richmond last play in a preliminary final"
q_emb = model.encode_query([query], convert_to_tensor=True)
score_vec = model.similarity(q_emb, corpus_emb)[0]
top_scores, top_idx = score_vec.topk(3)
for s, i in zip(top_scores.tolist(), top_idx.tolist()):
    print(f"{s:.4f}  {corpus[i][:100]}")

On a RTX 3090 the 4,874‑passage corpus encodes in ~20 s and each query scores in ~120 ms.

Retrieve‑and‑Rerank Pattern

Combine a fast dense bi‑encoder with a multi‑vector reranker to avoid building a full late‑interaction index:

from sentence_transformers import SentenceTransformer, MultiVectorEncoder, util

retriever = SentenceTransformer("jinaai/jina-embeddings-v5-text-nano-retrieval")
reranker  = MultiVectorEncoder("perplexity-ai/pplx-embed-v1-late-0.6b", trust_remote_code=True)

# 1️⃣ Retrieve top‑k with dense model
corpus_emb = retriever.encode_document(corpus, convert_to_tensor=True)
hits = util.semantic_search(retriever.encode_query([query], convert_to_tensor=True), corpus_emb, top_k=50)[0]
candidates = [corpus[h['corpus_id']] for h in hits]

# 2️⃣ Rescore candidates with MaxSim
q_emb = reranker.encode_query([query])
d_emb = reranker.encode_document(candidates)
rerank_scores = reranker.similarity(q_emb, d_emb)[0]
print(rerank_scores.argsort(descending=True)[:3])

Only the selected candidates are encoded as multi‑vectors, dramatically reducing memory while preserving late‑interaction quality.

Indexing Options

Several vector databases support native multi‑vector fields with MaxSim:

  • Qdrant (v1.10+)
  • Weaviate (v1.29+)
  • Vespa (long‑context ColBERT support)
  • LanceDB (v0.15.0+)
  • VectorChord (Postgres extension)
  • Milvus (v2.6.4, array‑of‑structs)
  • fast‑plaid (Rust implementation, approximate but fast)

Each system ingests the list of token‑level tensors returned by encode_document. Example for fast‑plaid (exact indexing):

from fast_plaid import search
fast_plaid = search.FastPlaid(index="nlp-index", device="cuda")
fast_plaid.create(documents_embeddings=document_emb)
results = fast_plaid.search(queries_embeddings=q_emb.unsqueeze(0), top_k=3)

The raw 608 k token vectors occupy 311 MB (float32) but compress to ~92 MB with PLAID’s centroid‑plus‑residual scheme.

Visual Document Retrieval

ColPali‑style models treat page images as token sequences of image patches. The same API works:

model = MultiVectorEncoder("vidore/colqwen2.5-v0.2")
queries = ["What is the variable on the y‑axis?", "Total outlay is maximum in which year?"]
images = [".../doc1.jpg", ".../doc2.jpg", ".../doc3.jpg", ".../doc4.jpg"]
q_emb = model.encode_query(queries)
d_emb = model.encode_document(images)
print(model.similarity(q_emb, d_emb))

A single page may yield hundreds of token vectors (e.g., 755 × 128), so token pooling becomes valuable (see next section).

Audio and Video Retrieval

The omni‑modal model vidore/colqwen-omni-v0.1 supports text, image, audio, and video without any transcription step.

model = MultiVectorEncoder("vidore/colqwen-omni-v0.1", model_kwargs={"dtype": torch.bfloat16})
# Audio example
audio = [...]  # list of raw waveforms (16 kHz mono)
q_emb = model.encode_query(["medicine for car nausea"])
d_emb = model.encode_document(audio, batch_size=2)
print(model.similarity(q_emb, d_emb)[0].topk(3))

Zero‑shot audio retrieval works because the model was trained on image‑text pairs and learns cross‑modal token alignments.

Interpretability

MaxSim’s per‑token decomposition enables exact attribution:

  • Heatmaps for images – overlay token‑level scores on page patches.
  • Text similarity maps – list the best‑matching document token for each query token.

The repository provides heatmap.py and text_similarity_map.py scripts that print token‑by‑token contributions and highlight the source passage.

Token Pooling to Reduce Index Size

HierarchicalTokenPooling clusters a document’s token vectors (Ward linkage on cosine distance) and replaces each cluster with its centroid, achieving roughly a 1 / pool_factor reduction.

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pool = HierarchicalTokenPooling(pool_factor=2)
pooled_emb = model.encode_document(docs, token_pooling=pool)

On the Natural Questions corpus, pool_factor=2 halves the index (311 MB → 156 MB) with only a ~0.4 % NDCG drop on BEIR. Higher factors give diminishing returns; evaluate on your own data before committing.

Speeding Up Inference

  • GPU – fp16 + Flash Attention (attn_implementation="flash_attention_2") yields ~2.4× throughput over fp32.
  • CPU – OpenVINO (when supported) and int8 quantization give modest speedups with <0.5 % accuracy loss.
  • Models that use query expansion masks (attend=False) cannot use Flash Attention because masked tokens would be dropped; use the default sdpa implementation instead.

Evaluation

MultiVectorNanoBEIREvaluator runs the 13‑subset NanoBEIR benchmark out‑of‑the‑box:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator
model = MultiVectorEncoder("lightonai/LateOn")
eval = MultiVectorNanoBEIREvaluator()
print(eval(model))

Results show LateOn (multi‑vector, 128‑dim) achieving a mean NDCG@10 of 0.6868, outperforming the dense counterpart lightonai/DenseOn (0.6764) on 9 of 13 datasets.

Migration from PyLate / colpali‑engine

PyLate Sentence Transformers
pylate.models.ColBERT(...) MultiVectorEncoder(...)
model.encode(..., is_query=True) model.encode_query(...)
model.encode(..., is_query=False) model.encode_document(...)
pylate.scores.colbert_scores model.similarity
pylate.indexes.PLAID keep using PyLate or switch to fast‑plaid / Qdrant
colpali‑engine Sentence Transformers
ColQwen2.from_pretrained(...) MultiVectorEncoder(...)
processor.process_queries(...) model.encode_query(...)
processor.process_images(...) model.encode_document(...)
processor.score_multi_vector(...) model.similarity(...)

Check the Migration Guide for details on prefixes, query expansion, and skip‑list handling.

Supported Models (excerpt)

Text Retrieval (selected)

Model Params Dim NanoBEIR
lightonai/LateOn-regularized 149 M 128 0.6897
lightonai/LateOn 149 M 128 0.6868
LiquidAI/LFM2.5-ColBERT-350M 353 M 128 0.6864
mixedbread-ai/mxbai-edge-colbert-v0-32m 32 M 64 0.6524
colbert-ir/colbertv2.0 110 M 128 0.6053

Visual Document Retrieval (selected)

Model Params Dim NanoViDoRe
webAI-Official/webAI-ColVec1.1-8b 8.4 B 640 0.6580
tencent/EVIE-Preview-4.5B 4.54 B 128 0.6405
vidore/colqwen2.5-v0.2 3.8 B 128 0.5402
vidore/colpali-v1.3 2.9 B 128 0.4802

All models listed on the Hub with the multi-vector tag are compatible; visual models may require a revision until their repository configuration PR is merged.

Acknowledgements

The implementation builds on ColBERT (Khattab & Zaharia, 2020), LightOn’s PyLate and fast‑plaid, the ColPali research team, and the token‑pooling work of Clavié, Chaffin, and Adams. Additional thanks to the MTEB benchmark contributors and all checkpoint authors.

Further Resources

  • Documentation – Usage, pretrained models, custom model creation, efficiency, and API reference (links in the original post).
  • Example scripts – Semantic search, retrieve‑and‑rerank, token pooling, heatmaps, and NanoBEIR evaluation.
  • Training guides – Overview, loss functions, and recipe scripts for LateOn/mLateOn.
  • Companion blogposts – Dense training, rerankers, sparse encoders, multimodal embeddings, and Matryoshka embeddings.

Sources

Related