Skip to content

Releases

v6.3.0 - Token-pooling-aware training and `kernels` support

Choose a tag to compare

@tomaarsen tomaarsen released this 11 Nov 14:45

Train multi-vector models for smaller document representations with token-pooling-aware and Matryoshka training. This release also adds optional kernels for MaxSim, Ward clustering, and transformer backbones, native AWS Neuron inference, faster utilities, and fixes across encoding, training, and evaluation.

Tip

I share much more about training multi-vector models with token-pooling-aware losses in my new Token Pooling-Aware Training for Multi-Vector Embedding Models blogpost. Give it a read!

Install with:

# Inference only, use one of:
pip install "sentence-transformers==6.3.0"
pip install "sentence-transformers[kernels]==6.3.0"
pip install "sentence-transformers[openvino]==6.3.0"

# Multimodal dependencies:
pip install "sentence-transformers[image]==6.3.0"
pip install "sentence-transformers[audio]==6.3.0"
pip install "sentence-transformers[video]==6.3.0"

# Training and merging support:
pip install "sentence-transformers[train]==6.3.0"
pip install "sentence-transformers[merge]==6.3.0"

# Or combine as needed:
pip install "sentence-transformers[train,image,kernels]==6.3.0"

Token-pooling-aware training (#4286)

Train for fewer document embeddings

Token pooling for multi-vector models merges similar document token embeddings to reduce the embeddings stored and scored during retrieval. With MultiVectorTokenPoolingLoss, you can now train the encoder for this compression so its document embeddings retain more useful information after pooling.

Wrap a retrieval loss with the pooling configurations you intend to deploy:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import (
    CachedMultiVectorMultipleNegativesRankingLoss,
    MultiVectorTokenPoolingLoss,
)
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("answerdotai/ModernBERT-base")
pooling = HierarchicalTokenPooling(pool_factor=4)
base_loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=64)
loss = MultiVectorTokenPoolingLoss(
    model,
    base_loss,
    token_poolings=[None, pooling],
    weights=[0.5, 0.5],
)

Pass loss to MultiVectorEncoderTrainer to train on the average of the unpooled and factor-4 retrieval losses. Because both variants reuse the same encoder embeddings, training adds clustering and scoring work without another encoder pass.

A matched experiment compared two ModernBERT-base models with fresh 128-dimensional projections, trained on the same 120,000 MS MARCO triplets with seed 42:

Training Unpooled nDCG@10 Factor-4 nDCG@10 Unpooled score retained
Ordinary 51.84 49.03 94.6%
Pooling-aware 52.39 51.23 97.8%

Factor-4 pooling reduces the average document representation from 193.22 to 48.68 embeddings, while the pooling-aware model retains 97.8% of its unpooled score.

To compress documents with the trained example model, pass a pooler to encode_document or set the runtime default with MultiVectorEncoder(..., token_pooling=...). By default, pooling leaves query embeddings unchanged:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling")
queries = model.encode_query(["Which planet is known as the Red Planet?"])
documents = model.encode_document(
    ["Mars has a reddish surface.", "Jupiter is the largest planet."],
    token_pooling=HierarchicalTokenPooling(pool_factor=4),
)
scores = model.similarity(queries, documents)

See the blogpost for the method and experiments, and the training guide for complete scripts. (#4286)

Adaptive token budgets

Choose one of three HierarchicalTokenPooling controls to match the document budget you want to deploy:

Control Effect
pool_factor=4 Keep approximately one quarter of each document's vectors.
max_tokens=16 Keep at most 16 vectors per document, including protected tokens.
distance_threshold=... Allocate vectors according to embedding geometry using a shared Ward distance threshold.

For an adaptive budget, calibrate a threshold from representative document embeddings so the pooler targets an average count while allocating vectors according to each document's content:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling")
calibration_documents = [
    "Mars is the fourth planet from the Sun and has a reddish surface.",
    "Jupiter is a gas giant with many moons and a prominent Great Red Spot.",
]
embeddings = model.encode_document(calibration_documents)
threshold = HierarchicalTokenPooling.calibrate_distance_threshold(embeddings, target_tokens=8)
pooling = HierarchicalTokenPooling(distance_threshold=threshold) if threshold is not None else None
compressed = model.encode_document(calibration_documents, token_pooling=pooling)

The threshold targets an average on the calibration sample, so individual documents can receive more or fewer vectors and another corpus can have a different average. If the sample already fits the budget, calibration returns None and no pooling is needed.

These controls also work without pooling-aware training, so you can try compression on an existing model and evaluate it by passing the poolers to multi-vector evaluators.

Matryoshka training for multi-vector models

MultiVectorMatryoshkaLoss trains a model to retain retrieval quality when token embeddings are truncated. Wrap your retrieval loss with the dimensions you want to support:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.losses import (
    CachedMultiVectorMultipleNegativesRankingLoss,
    MultiVectorMatryoshkaLoss,
)

model = MultiVectorEncoder("answerdotai/ModernBERT-base")
base_loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=64)
loss = MultiVectorMatryoshkaLoss(
    model,
    base_loss,
    matryoshka_dims=[128, 64, 32],
    matryoshka_weights=[1 / 3] * 3,
)

After training, use the shorter vectors by passing truncate_dim=64 to both encode_query and encode_document, or set that width as the model constructor default.

In a separate matched experiment with 30,000 MS MARCO triplets, Matryoshka training preserves more retrieval quality at 64 dimensions, measured by mean nDCG@10 (%) across 13 NanoBEIR datasets:

Training 128-dimensional nDCG@10 64-dimensional nDCG@10 Full-width score retained
Ordinary 48.45 45.46 93.8%
Matryoshka 48.95 47.11 96.2%

Using 64 dimensions halves vector storage at unchanged precision, with retention measured against each model's own 128-dimensional score.

To train for fewer vectors and shorter dimensions together, wrap MultiVectorMatryoshkaLoss with MultiVectorTokenPoolingLoss so pooling happens before truncation. See the Matryoshka guide for complete training and evaluation scripts.

Optional kernels for MaxSim, token pooling, and transformer backbones (#4344)

With use_kernels=True, all four model types can use compatible Hub kernels to accelerate operations in their transformer backbones. For multi-vector models, the same setting also enables kernels for MaxSim and MeanMaxSim scoring, as well as the Ward clustering used in token pooling:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

model = MultiVectorEncoder("lightonai/mLateOn", use_kernels=True)
queries = model.encode_query(["How can I reduce retrieval storage?"])
documents = model.encode_document(
    ["Token pooling merges similar document vectors before indexing."],
    token_pooling=HierarchicalTokenPooling(pool_factor=2),
)
scores = model.similarity(queries, documents)

Install the optional dependencies with pip install "sentence-transformers[kernels]" before enabling kernels. The default multi-vector providers target CUDA and require Triton for MaxSim acceleration, falling back to reference implementations when a compatible device or provider build is unavailable.

In the Ward benchmark, the kernel reduces pooling time from 7.00 seconds with fastcluster to 1.02 seconds, about 6.9x faster, for 5,000 SciFact titles and abstracts. Including encoding, the improvement is smaller: 23.14 to 18.91 seconds, an 18% reduction with mLateOn, BF16, batch size 32, factor 2, and an RTX 3090.

Kernels are opt-in runtime settings, so enable them again when reloading a saved model. For finer control, pass a Transformers KernelConfig to choose providers separately for the backbone, scoring, and pooling operations, as described in the kernel guide. (#4344)

Native AWS Neuron inference (#4293)

After installing TorchNeuron in a supported AWS environment, set device="neuron" to run dense, cross-encoder, and multi-vector inference through the existing torch backend:

from sentence_transformers import Sente...
Read more

v6.2.0 - Model merging, faster loading, and distributed evaluation

Choose a tag to compare

@tomaarsen tomaarsen released this 21 Oct 12:27

This release adds model merging for all four model types, reduces redundant Hub requests when loading models, and distributes evaluator inference across DDP ranks. It also includes lower-memory training losses, more flexible multi-vector compression, and fixes throughout inference, training, and evaluation.

Install this version with

# Inference only, use one of:
pip install sentence-transformers==6.2.0
pip install sentence-transformers[openvino]==6.2.0

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==6.2.0
pip install sentence-transformers[audio]==6.2.0
pip install sentence-transformers[video]==6.2.0

# Training + Inference
pip install sentence-transformers[train]==6.2.0
pip install sentence-transformers[merge]==6.2.0

# Or combine as needed:
pip install sentence-transformers[train,image]==6.2.0

Model merging for all four model types (#3776)

Combine compatible fine-tuned checkpoints without another training run using SentenceTransformer.merge, SparseEncoder.merge, CrossEncoder.merge, or MultiVectorEncoder.merge. The integration uses Arcee AI's mergekit and supports well over a dozen methods such as Linear, SLERP, TIES, and DARE.

Install the merge extra with pip install "sentence-transformers[merge]". For example, take a weighted average of two compatible MiniLM models:

from sentence_transformers import SentenceTransformer

models = [
    "sentence-transformers/all-MiniLM-L6-v2",
    "sentence-transformers/multi-qa-MiniLM-L6-cos-v1",
]

merged = SentenceTransformer.merge(
    models,
    method="linear",
    parameters={"weight": [0.6, 0.4]},
)
merged.save_pretrained("merged-MiniLM-L6")

embeddings = merged.encode(["Combining compatible model weights."])

Or interpolate between the same inputs with SLERP:

merged = SentenceTransformer.merge(
    models,
    method="slerp",
    parameters={"t": 0.5},
)

The method returns a model that can be used, trained, and saved normally. Configuration, prompts, and processors come from the first input. Inputs must have compatible architectures, module configurations, and tokenizers. Merging uses FP32 on CPU by default and holds the inputs and result in memory.

See the merging guides for Sentence Transformers, Cross Encoders, Sparse Encoders (plus an example evaluation!), and Multi-Vector Encoders. Be sure to evaluate the merged model on your own tasks, since combining checkpoints can improve some tasks while regressing others.

Faster Hub model loading (#3890)

Combined with huggingface_hub>=1.26.0, model loading now resolves the requested branch or tag to a commit once and reuses that result across component downloads, reducing repeated Hub requests. Older Hub versions retain the previous loading behavior.

The loading benchmarks in #3890 measured constructor time in fresh processes with a warm cache and online Hub access. Values are medians of seven alternating runs per configuration:

Model Before After HTTP requests before / after
sentence-transformers-testing/stsb-bert-tiny-safetensors 5.78 s 2.26 s 34 / 8
sentence-transformers/all-MiniLM-L6-v2 6.39 s 2.66 s 35 / 9
sparse-encoder-testing/splade-bert-tiny-nq 5.77 s 2.25 s 34 / 8
cross-encoder/ms-marco-MiniLM-L6-v2 4.38 s 2.48 s 23 / 9

Constructor time fell by 43 to 61 percent in these workloads. Loading models from local directories is unchanged and unaffected by this improvement.

Distributed evaluator inference during DDP training (#4004)

During DDP training, evaluators now share their encode() and predict() work across the existing training processes for all four model types, with each rank processing a shard of the inputs. Rank zero combines predictions in their original order, computes metrics over the full dataset, and writes evaluation output once. Metrics are then broadcast to every rank for checkpoint selection and early stopping.

This now happens automatically when using an evaluator through the trainer. It also covers encode_query(), encode_document(), and rank(), which use those inference methods internally.

Explicit device and pool arguments retain their usual behavior, and calls with arguments that cannot be serialized fall back to rank zero. Predictions are gathered through CPU memory, so the speedup depends on inference cost and communication overhead. Evaluator support for FSDP and DeepSpeed is unchanged.

Smaller Features

Experimental FP32 layers for CrossEncoders (#4230)

CrossEncoders can compute selected linear layers in FP32 to reduce BF16 score ties while keeping a reduced-precision backbone:

from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L6-v2",
    device="cuda",
    model_kwargs={"dtype": "bfloat16"},
    keep_in_fp32_modules=["0.model.classifier"],
)

This uses the exact model.named_modules() paths. The selection is saved with the model, works under autocast during training and inference, and leaves stored parameter dtypes unchanged. In my experiments, this doesn't improve NDCG@10, but solely helps reduce ties caused by the low BF16/FP16 output precision.

Lower memory training losses (#4026, #4029, #4030)

  • BatchSemiHardTripletLoss uses quadratic intermediate memory instead of constructing cubic distance and mask tensors. It also corrects a gradient issue when selected negative distances tie (#4026).
  • BatchAllTripletLoss removes invalid anchor-positive pairs before constructing triplets. Memory is quadratic for paired labels, though large label groups can still require cubic storage (#4029).
  • CoSENTLoss, AnglELoss, and their sparse counterparts use sorted suffix reductions with linear intermediate storage. MPS execution requires PyTorch 2.9 or newer (#4030).

Multi-vector compression and scoring (#3965, #3941, #4008, #4153)

  • HierarchicalTokenPooling accepts fractional pool_factor values for finer control over compression (#3965) and can use optional fastcluster linkage for faster clustering (#3941).
  • HierarchicalTokenPooling now normalizes output tokens by default, including protected tokens. Older saved configurations retain their previous behavior. Set normalize_embeddings=False on the pooling module to opt out (#4008).
  • MaxSim scoring budgets score blocks using both query and document token counts, reducing memory use for long query batches. A single query-document pair can still exceed chunk_elements (#4153).

Compatibility Notes

  • Python and Hub versions: Added Python 3.14 and huggingface_hub 2.x support, retaining Python 3.10 support and the Hub 1.3.0 minimum (#4162, #4003).
  • Scalar quantization: Inferring int8 or uint8 calibration ranges from a single vector now raises ValueError. Encode floats first, then call quantize_embeddings with representative calibration_embeddings or explicit ranges (#4150).
  • Checkpoint paths: Module paths must be relative and cannot contain .. path components, even with trust_remote_code=True. Keep custom modules inside the checkpoint tree (#4175).

Bug Fixes

Pooling precision and normalization

  • Improved FP16/BF16 mean and weighted-mean pooling arithmetic, corrected weighted positions under left padding, and fixed prompt exclusion in flattened CLS fallback selection (#4106, #4104, #4159).
  • Corrected normalization of single sparse vectors and sparse matrices with non-default dtypes or low-precision values (#4149, #4045).

Multi-process inference

  • Worker exceptions are reported to the caller, and terminated workers are detected while polling for results (#3920, #4146).
  • Worker tensor outputs are moved to CPU before queueing, including tensors nested in lists and dictionaries. Hard-negative mining also stops its pool if mining fails (#3971, #4086).

Model placement and memory

  • Improved device handling for quantized models and models using Accelerate device maps or offloading, and reduced repeated placement work. Call `clear_pl...
Read more

v6.1.0 - Documentation and input handling improvements

Choose a tag to compare

@tomaarsen tomaarsen released this 18 Sep 10:44

This is a small release with refreshed inference benchmarks, clearer documentation, and a few improvements to multimodal input handling.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==6.1.0

# Inference only, use one of:
pip install sentence-transformers==6.1.0
pip install sentence-transformers[onnx-gpu]==6.1.0
pip install sentence-transformers[onnx]==6.1.0
pip install sentence-transformers[openvino]==6.1.0

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==6.1.0
pip install sentence-transformers[audio]==6.1.0
pip install sentence-transformers[video]==6.1.0

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==6.1.0

Changes

  • Refreshed the inference benchmarks and efficiency guide, clarified benchmark labels, and documented ONNX CPU thread tuning (#3977, #4007, #3987).
  • Added a shared input formats guide and SPARSEUP to the pretrained sparse model tables (#4036, #4038).
  • Improved multimodal input handling and chat-template argument routing (#4037, #4048).
  • Corrected loss formulas and argument descriptions in docstrings, clarified custom-module import requirements, and fixed links, typos, and optional type annotations.

For existing multimodal dict inputs, message content now follows each input's key order, which can change embeddings or scores. Video-frame URLs without recognizable image extensions need an explicit frame wrapper to remain one video. See #4037 for details.

All Changes

  • [docs] Document that optional imports in a custom module become load-time requirements by @abtonmoy in #3922
  • [docs] Name the metric on the quality axis of the backend benchmark figures by @tomaarsen in #3977
  • docs: fix Constrastive/heaqp typos and subject-verb agreement by @simpleqt in #3983
  • docs: drop stale margin/squared Args from triplet loss docstrings by @simpleqt in #3981
  • docs: fix cross-encoder link and 'show case' wording in examples READMEs by @simpleqt in #3982
  • Correct the inverted CoSENT loss formula in two docstrings by @vineethsaivs in #3984
  • Fix implicit Optional annotations by @VenishPaneliya in #3990
  • [docs] Clarify sparse COO inputs for semantic_search by @tomaarsen in #4001
  • [docs] Refresh inference benchmarks and promote efficiency guides by @tomaarsen in #4007
  • [docs] Describe the MarginMSE target as the signed margin by @MaxFreedomPollard in #4009
  • Document ONNX CPU thread tuning and benchmark settings by @ErnisBadmae in #3987
  • docs: Document shared input formats across model archetypes by @tomaarsen in #4036
  • [docs] Add SPARSEUP to pretrained sparse models by @tomaarsen in #4038
  • Support multiple items per modality in one input by @tomaarsen in #4037
  • Fix load_audio_from_video routing in chat templates by @tomaarsen in #4048

Full Changelog: v6.0.1...v6.1.0

v6.0.1 - Restore the PyLate prefix on prompted checkpoints, 80 documented multi-vector models

Choose a tag to compare

@tomaarsen tomaarsen released this 31 Aug 07:50

This patch release fixes a multi-vector loading bug: PyLate checkpoints that carry both a [Q]/[D] prefix and a text prompt lost the prefix, so they were encoded without a marker they were trained with. It also grows the documented multi-vector model tables from 51 checkpoints to 80.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==6.0.1

# Inference only, use one of:
pip install sentence-transformers==6.0.1
pip install sentence-transformers[onnx-gpu]==6.0.1
pip install sentence-transformers[onnx]==6.0.1
pip install sentence-transformers[openvino]==6.0.1

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==6.0.1
pip install sentence-transformers[audio]==6.0.1
pip install sentence-transformers[video]==6.0.1

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==6.0.1

Compose the PyLate prefix with the saved prompt (#3968)

PyLate supports a prefix as well as prompt text, and Sentence Transformers assumed that these were mutually exclusive. However, the following three checkpoints carry both a prefix and a prompt, and were trained with the prefix prepended to the prompt text:

PyLate trained these checkpoints on [CLS] [Q] search_query: ... and [CLS] [D] search_document: ..., but Sentence Transformers dropped the prefix and encoded them as [CLS] search_query: ... and [CLS] search_document: .... After the fix, the performance of the ColBERT-Zero model improved from 0.6569 NDCG@10 to 0.6824 on the NanoBEIR benchmark. To my knowledge, only these 3 checkpoints used both a prefix and a prompt, so this is the only case where the bug would have affected you.

Forward task to routed modules in Router.preprocess (#3967)

Router.forward has always passed task down to the routed module, but Router.preprocess did not, so anything a module does with the task at preprocessing time did nothing behind a Router: query_length and document_length caps, query_expansion, and the chat-template task keyword. No released checkpoint combines a Router with those settings, so this is a latent bug rather than one you are likely to have hit. It would have affected anyone building such a model themselves, with no error to indicate it.

Documentation

  • The multi-vector pretrained models tables grew from 29 text and 22 visual document retrieval checkpoints to 38 and 42 (#3963, #3969), with revision and trust_remote_code notes refreshed as upstream pull requests merged (#3972).
  • Documented which Hub tag to filter on for each model type (#3964), and linked the Multi-Vector Encoder blogposts from the docs (#3966).
  • Added MultiVectorEncoder to the Agent Skill README and refreshed the SauerkrautLM scores (#3955).
  • Corrected the minimum versions in the README (#3952). It still recommended PyTorch 1.11.0+ and transformers v4.41.0+, where v6.0 requires PyTorch 2.2+ and transformers v5.0+. This is the text rendered as the PyPI project description.

All Changes

  • docs: Add MultiVectorEncoder to the skills README, refresh the SauerkrautLM scores by @tomaarsen in #3955
  • Fix the README dependency versions by @tomaarsen in #3952
  • [docs] Add more MultiVectorEncoder models to the pretrained models tables by @tomaarsen in #3963
  • [docs] Clarify which Hub tags to filter on for each model type by @tomaarsen in #3964
  • [docs] Link the Multi-Vector Encoder blogposts by @tomaarsen in #3966
  • [fix] Forward task to routed modules in Router.preprocess by @tomaarsen in #3967
  • [fix] Compose the PyLate prefix with the saved prompt by @tomaarsen in #3968
  • [docs] Refresh the MultiVectorEncoder pretrained models tables by @tomaarsen in #3969
  • Add/update revisions for pretrained MVE models by @tomaarsen in #3972

Full Changelog: v6.0.0...v6.0.1

v6.0.0 - MultiVectorEncoder for ColBERT & late interaction models, transformers v5, float32 scoring, faster training & encoding

Choose a tag to compare

@tomaarsen tomaarsen released this 18 Aug 13:46

This major release introduces Multi-Vector Embedding models, also known as late interaction or ColBERT-style models, as a fourth model type alongside SentenceTransformer, CrossEncoder, and SparseEncoder. Going forward, you'll be able to use Sentence Transformers for training, inferencing, and interpreting Multi-Vector Embedding models.

It also modernizes the dependency floors to transformers v5, fixes a class of silent scoring bugs caused by half precision, and speeds up both training and encoding.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==6.0.0

# Inference only, use one of:
pip install sentence-transformers==6.0.0
pip install sentence-transformers[onnx-gpu]==6.0.0
pip install sentence-transformers[onnx]==6.0.0
pip install sentence-transformers[openvino]==6.0.0

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==6.0.0
pip install sentence-transformers[audio]==6.0.0
pip install sentence-transformers[video]==6.0.0

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==6.0.0

Tip

Our Multi-Vector (Late Interaction) Embedding Models with Sentence Transformers blogpost is an excellent place to learn about multi-vector models: loading the various checkpoint formats, encoding and scoring, plugging them into a search stack, running them on page images, and keeping the index affordable.

Warning

This is a major release with breaking changes. Upgrading from v5.x to v6.0 may require code updates. The changes marked 🚨 below are the ones most likely to affect you, and the Migration Guide has the full list. If you run into issues when upgrading, feel free to open an issue.

MultiVectorEncoder: ColBERT-style late interaction models (#3794)

Sentence Transformers v6.0 introduces MultiVectorEncoder, for ColBERT-style late interaction retrieval. Where a regular embedding model compresses a whole text into one vector, a multi-vector model keeps one vector per token and scores query against document with the MaxSim operator. That preserves token-level matching information that a single vector has to average away, which usually means stronger retrieval at the cost of a bigger index. It is also the state of the art for visual document retrieval, where a text query is matched against page images directly, with no OCR step in between.

Any PyLate checkpoint and any Stanford-NLP ColBERT checkpoint loads straight into it, and colpali-engine models for visual document retrieval work too, through the same familiar API you already use for dense, sparse, and reranker models.

from sentence_transformers import MultiVectorEncoder

# Download from the 🤗 Hub
model = MultiVectorEncoder("lightonai/LateOn")

query_embeddings = model.encode_query(["Which planet is known as the Red Planet?"])
document_embeddings = model.encode_document([
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot.",
    "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
])

print(query_embeddings[0].shape)
# (12, 128)

scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[10.7942, 11.1104, 10.9743, 11.0811]])

Mars wins, as it should, though notice how close the four scores are. That is normal for MaxSim: the scores often look similar, but the ranking is still exact. The blogpost explores this in more detail.

Note what you get back: a list of 2D tensors on the model device, one per input, each of shape (num_tokens, embedding_dim). Unlike dense embeddings, you cannot stack these into one rectangular tensor, because every input has its own token count. Pass convert_to_numpy=True for a list of numpy arrays instead, which is what you want once a corpus outgrows device memory.

Multi-vector models are also asymmetric: queries and documents go through different prefixes, different length caps, and different scoring masks. Unlike many dense models, where the two are interchangeable, encode_query and encode_document are required to get correct embeddings.

The MaxSim operator

Scoring uses MaxSim: for each query token, take its highest similarity against any document token, then sum those maxima across the query.

$$\text{MaxSim}(Q, D) = \sum_{Q_i \in Q} \max_{D_j \in D} Q_i \cdot D_j$$

You can read the operator as a soft alignment: every query token points at the one document token that best explains it, and the score is how well the document explains the query overall. The alignment does not have to be lexical, since the token embeddings are contextualized. But when an exact match does matter to you (a product code, a surname, a function name), MaxSim has a token sitting right there to match it, where a single-vector model had to fold it into an average.

Because MaxSim sums over query tokens, its magnitude scales with the query token count, so scores are not comparable across models with different query recipes. If you want scores on a bounded scale, use similarity_fn_name="meanmaxsim", which divides by the query token count and gives you an average cosine similarity in [-1, 1].

Scoring builds a 4-dimensional intermediate of every query token against every document token, which is the largest tensor in the operation. Every scoring function takes a chunk_elements budget that bounds it, defaulting to 100 million elements (roughly 400 MB in float32), so lower it if you run out of memory. Scores and gradients are bit-identical whatever you set it to. maxsim and maxsim_pairwise also take a device, which scores one chunk at a time on that device and moves each result straight back, letting you score a corpus larger than your VRAM on the GPU. Both are reachable through similarity, which forwards any extra keyword arguments to the scoring function:

scores = model.similarity(query_embeddings, document_embeddings, chunk_elements=1_000_000, device="cuda")

When training, pass the budget to the loss instead, with similarity_fct=partial(colbert_scores, chunk_elements=1_000_000). It chunks the document axis, so it composes with the loss-level score_mini_batch_size, which chunks the query axis.

Are they any good?

lightonai/LateOn and lightonai/DenseOn were trained by LightOn on the same data with the same ModernBERT backbone and the same 149M parameters, differing only in whether they keep one vector per token or pool down to one per document. Running both over all 13 NanoBEIR datasets isolates what that choice buys:

NanoBEIR dataset LateOn (multi-vector, 128d) DenseOn (dense, 768d)
MSMARCO 0.7194 0.6517
NQ 0.7810 0.7511
HotpotQA 0.9295 0.8802
FEVER 0.9702 0.9612
ClimateFEVER 0.4887 0.4846
DBPedia 0.6836 0.6748
QuoraRetrieval 0.9795 0.9687
Touche2020 0.5938 0.5673
ArguAna 0.5562 0.5660
NFCorpus 0.3949 0.3851
SciFact 0.7978 0.8057
SCIDOCS 0.4469 0.4484
FiQA2018 0.5871 0.6491
Mean 0.6868 0.6764

Late interaction wins on 9 of the 13 datasets and on the mean, by roughly one NDCG point. The four it loses (ArguAna, FiQA2018, SCIDOCS, and SciFact) are the shape of the tradeoff you should expect: a real gain in retrieval quality at the same model size, paid for in index footprint, rather than a universal win on every dataset. The same pair scores 57.22 against 56.20 on the full 15-dataset BEIR, a comparable gap, so the margin is not an artifact of the small benchmark.

That footprint is the real cost. One vector per token instead of one vector per document is a lot more vectors, only partly offset by the smaller dimension. Encoding 4,874 Natural Questions passages with lightonai/LateOn produced 608,414 token vectors, an average of 124.8 per passage:

Representation Vectors Dimensions float32 index
Dense, all-MiniLM-L6-v2 4,874 384 7.5 MB
Dense, gte-modernbert-base 4,874 768 15.0 MB
Multi-vector, LateOn 608,414 128 311.5 MB

That is about 42x the storage of the MiniLM index. Token Pooling cuts the vector count before any of that, real late interaction indexes compress heavily (the same vectors take 88 MB as a fast-plaid PLAID index), and using a multi-vector model as a reranker over a dense first stage avoids building an index at all.

Every checkpoint format loads

Multi-vector checkpoints have been published in several formats over the years. MultiVectorEncoder reads all of them, so loading looks the same whatever the model started life as:

from sentence_transformers import MultiVectorEncoder

# Native Sentence Transformers checkpoints. PyLate builds on the same schema,
# so any PyLate checkpoint loads identically
model = MultiVectorEncoder("lightonai/LateOn")
model = MultiVectorEncoder("mixedbread-ai/mxbai-edge-colbert-v0-17m")
model = MultiVectorEncoder("LiquidAI/LFM2-ColBERT-350M")

# Any Stanford-NLP ColBERT checkpoint, detected via the `HF_ColBERT` architecture
# marker. The inline projection weight and the recipe come fro...
Read more

v5.7.0 - GradCache Overhaul, torch.compile Inference, and a Large Batch of Correctness Fixes

Choose a tag to compare

@tomaarsen tomaarsen released this 06 Aug 12:12

This minor version is a correctness and performance-focused release. It rebuilds all gradient-cached losses on one shared engine, fixing several silently wrong gradients and adding token-based mini-batching for up to 3.9x faster cached-loss training. It also makes model.compile() actually speed up inference, and brings a long list of fixes across embedding quantization, evaluators, hard-negative mining, community detection, and multimodal inputs.

Two changes are marked breaking (🚨): int8/uint8 embedding quantization now clips out-of-range values and floors bucket values, so int8 outputs are no longer bit-identical with earlier versions, and AdaptiveLayerLoss/Matryoshka2dLoss now weight prior-layer losses uniformly by default. There's also a forward-looking deprecation: loading models whose modules import classes from outside sentence_transformers will require trust_remote_code=True from v6.0.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==5.7.0

# Inference only, use one of:
pip install sentence-transformers==5.7.0
pip install sentence-transformers[onnx-gpu]==5.7.0
pip install sentence-transformers[onnx]==5.7.0
pip install sentence-transformers[openvino]==5.7.0

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==5.7.0
pip install sentence-transformers[audio]==5.7.0
pip install sentence-transformers[video]==5.7.0

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==5.7.0

GradCache overhaul: correct gradients and token-based mini-batching (#3862)

The gradient-cached losses (CachedMultipleNegativesRankingLoss, CachedGISTEmbedLoss, CachedSpladeLoss, the Cross Encoder CachedMultipleNegativesRankingLoss, and MegaBatchMarginLoss) train with large batch sizes at constant memory by embedding in mini-batches and replaying them with cached gradients. Each loss carried its own diverged copy of that machinery. They are now all rebuilt on one shared engine, which fixed several bugs that silently corrupted gradients:

  • Cross Encoder CachedMultipleNegativesRankingLoss on GPU: the backward pass used different dropout masks than the forward pass, silently biasing gradients for every reranker trained with dropout active on CUDA or MPS. CPU training was unaffected.
  • Running two forward passes before a backward pass (e.g. in custom training loops) made CachedGISTEmbedLoss and the Cross Encoder loss backpropagate the wrong batch's gradients, because the cache was stored on the loss module. The cache now travels with each forward pass's backward hook (the .cache and .random_states loss attributes are gone as a result).
  • Pooling(include_prompt=False) (e.g. Instructor models) mutated the attention mask in place, so the backward re-embedding of every cached loss ran with a different mask than the forward pass.
  • MatryoshkaLoss(GISTEmbedLoss(...)): the guide model overwrote the cached embeddings, so only the largest Matryoshka dimension was actually trained.

Along the way, this also fixed an autocast dtype crash in the backward pass and the trainer retaining autograd graphs between logging steps when tracking loss components.

MegaBatchMarginLoss's default mini-batched version is rebuilt on the engine as well. It crashed outright on recent releases, and underneath that, its historical implementation only applied the last mini-batch's gradients. It now trains on the full batch (results will differ, for the better), works with MatryoshkaLoss, evaluates under torch.no_grad, and raises for a third input column instead of silently ignoring it.

The headline feature is mini_batch_num_tokens, available on CachedMultipleNegativesRankingLoss, CachedMultipleNegativesSymmetricRankingLoss, CachedGISTEmbedLoss, CachedSpladeLoss, and MegaBatchMarginLoss. Instead of a fixed number of sequences per mini-batch, mini-batches are greedily packed by total non-padding token count, giving near-constant work per mini-batch on variable-length data:

from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.losses import CachedMultipleNegativesRankingLoss

model = SentenceTransformer("microsoft/mpnet-base")
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_num_tokens=16384)

On the PR's Natural Questions benchmark, cached-loss training with flash attention and a tuned token budget dropped from 715 to 182 seconds (3.9x) versus the previous release, with unchanged quality. The engine also trims trailing padding from each mini-batch, which alone is worth about 26% throughput on the default padded path. The updated training efficiency documentation recommends the smallest token budget that saturates your GPU. mini_batch_size keeps working everywhere as before.

model.compile() now speeds up encode() and predict() (#3848)

encode() and predict() previously called the model's forward() directly, bypassing nn.Module.__call__, which is where torch.compile installs its compiled path. As a result, model.compile() was silently a no-op for inference. The forward pass now runs through __call__, so compilation applies (outputs are bit-identical when not compiling).

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-small-en-v1.5", model_kwargs={"torch_dtype": "bfloat16"})
model.compile(dynamic=True)

# Compilation is lazy, so warm up on representative inputs before benchmarking or serving
embeddings = model.encode(["This is an example sentence", "Each sentence is converted"])

For the largest gains at batch size 1, mode="reduce-overhead" applies CUDA graphs: measured in bf16 on an RTX 3090, roughly 3.0x on bge-small-en-v1.5 and 3.7x on modernbert-embed-large. Very small models like all-MiniLM-L6-v2 see little gain or even a slowdown, as their inference is dominated by tokenization and Python overhead, so always measure on your own model and hardware. The Speeding up Inference documentation for all three model types (SentenceTransformer, CrossEncoder, SparseEncoder) gains a torch.compile tab covering dynamic shapes, CUDA graphs, and the fixed-shape padding the latter need.

🚨 Embedding quantization fixes (#3847, #3825)

Two fixes for quantize_embeddings:

  • 🚨 Clip and floor in int8/uint8 quantization (#3847): values outside the calibration range used to wrap around in the float-to-integer cast, e.g. producing [-42, 42, -112] where [-128, 42, 127] was correct. Out-of-range values now saturate at the bounds. Bucket values are also floored before casting, which makes int8 a consistent uniform quantizer matching uint8, but shifts in-range int8 outputs in the lower half of the range down by one level out of 256. int8 quantization is therefore not bit-identical with earlier versions: re-quantize existing int8 corpora rather than mixing old and new quantized embeddings. uint8 outputs only change where they previously wrapped. Resolves #3159.
  • Pack bits per embedding in binary/ubinary quantization (#3825): precision="binary"/"ubinary" crashed with a ValueError when the embedding dimension was not a multiple of 8 (e.g. some Matryoshka-truncated dimensions), because the whole batch was packed as one flat bit array. Bits are now packed per embedding via np.packbits(..., axis=-1). Output is bit-identical for dimensions divisible by 8, like 384, 768, or 1024.

🚨 AdaptiveLayerLoss: detached KL teacher and configurable layer weighting (#3880, #3901)

AdaptiveLayerLoss and Matryoshka2dLoss train models whose embeddings remain useful when transformer layers are dropped, by also training each prior layer's output, including a KL-divergence term that distills the final layer's embeddings into the prior layers. Two changes, together resolving #3757:

  • Detach the KL teacher (#3880): the final-layer embeddings used as the KL target were not detached, so the KL gradients also flowed back into the final layer, pulling the teacher toward its own students. The teacher is now a standard stop-gradient target.
  • 🚨 Configurable layer_weighting (#3901): the prior-layer losses were hardcoded to decay as 1 / (1 + layer_idx). The new layer_weighting option accepts "uniform" (the new default, matching the 2DMSE and Starbucks papers), "log" (the ESE paper), "linear" (exactly the previous behavior), or any callable mapping a layer index to a weight. Uniform outperformed log and linear at every truncated layer count on STSb, hence the new default. Pass layer_weighting="linear" to restore the previous behavior exactly.

Clear errors for LLM-based rerankers and embedders (#3902, #3869)

Two silent failure modes of decoder-based models now raise actionable errors:

  • Chat templates that cannot carry query/document pairs (#3902): turning a general instruct LLM into a CrossEncoder reranker previously failed silently, as such a chat template does not know the query and document roles: both messages rendered to nothing, producing a cryptic IndexError or meaningless scores from an empty prompt. The first predict() call or training batch containing a pair now raises an error that names the model, explains why the roles are dropped, and includes a copyable Jinja template remedy. Published rerankers with working templates are unaffected (the 20 most-downloaded chat-template rerankers on the Hub all pass). Resolves #3881.
  • Causal inputs that are not left-padded (#3869): models that score or pool from the final token position require left padding. The previous attribute-based check never fired for multimodal p...
Read more

v5.6.1 - Flash Attention Fix for XLM-R and RoBERTa Models

Choose a tag to compare

@tomaarsen tomaarsen released this 23 Jul 14:42

v5.6.1 - Flash Attention Fix for XLM-R and RoBERTa Models

This patch release fixes silently degraded embeddings for RoBERTa-family models when flash attention is requested with transformers v5, notably every XLM-R based multilingual embedding model (BAAI/bge-m3, intfloat/multilingual-e5-large, etc.). The bug affected v5.5.0, v5.5.1, and v5.6.0.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==5.6.1

# Inference only, use one of:
pip install sentence-transformers==5.6.1
pip install sentence-transformers[onnx-gpu]==5.6.1
pip install sentence-transformers[onnx]==5.6.1
pip install sentence-transformers[openvino]==5.6.1

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==5.6.1
pip install sentence-transformers[audio]==5.6.1
pip install sentence-transformers[video]==5.6.1

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==5.6.1

Fix position_ids offsetting for RoBERTa-family models when flattening inputs (#3879)

Since v5.5.0, the Transformer module flattens text-only batches into one packed sequence when flash attention is requested, skipping all padding overhead for a notable performance improvement. The position_ids of that packed sequence restart at 0 for every text, which is correct for the vast majority of models. RoBERTa-family architectures however compute positions as padding_idx + 1 + n for the n-th token, so every token read a position embedding shifted by padding_idx + 1 (usually 2). Nothing crashes, the embeddings are just silently worse.

from sentence_transformers import SentenceTransformer

# An affected configuration: flash attention with an XLM-R based model
model = SentenceTransformer(
    "BAAI/bge-m3",
    model_kwargs={"attn_implementation": "flash_attention_2"},
)

Measured on BAAI/bge-m3:

Evaluation padded packed, 0-based positions packed, with this fix
stsb test Spearman 0.8485 0.7239 0.8485
NanoBEIR mean nDCG@10 0.6041 0.5414 0.6050

The quality loss recovers exactly once the offset is applied. The fix scans the loaded model's modules once for an int padding_idx stored next to a learned position_embeddings table, and offsets the packed position_ids when that pair is found. An audit of transformers finds 16 architectures with that pair (roberta, xlm_roberta, xlm_roberta_xl, camembert, roberta_prelayernorm, xmod, data2vec_text, longformer, luke, ibert, mpnet, markuplm, lilt, layoutlmv3, esm, and pp_doclayout_v2), all offset by exactly padding_idx + 1, and no 0-based or rotary architecture matches.

You are only affected if you encoded text with flash attention requested on transformers v5 with a RoBERTa-family checkpoint. The default padded path (e.g. sdpa) was never affected, and neither were MPNet models like all-mpnet-base-v2 despite mpnet appearing in the audit: transformers does not support flash attention for MPNet at all. If you did index a corpus with such a configuration, re-encode it after upgrading: pre-fix embeddings score notably worse and do not mix with post-fix embeddings.

What's Changed

  • [ci] Exclude librosa/numba/llvmlite on Python 3.13 by @tomaarsen in #3835
  • [tests] Skip bf16 + Windows + CPU forwards, as they can WindowsError on torch 2.13 by @tomaarsen in #3863
  • 🚨 Fix position_ids offsetting for RoBERTa-family models when flattening inputs by @tomaarsen in #3879

Full Changelog: v5.6.0...v5.6.1

v5.6.0 - Fixes for Causal LM Rerankers, Hard-Negative Mining, and More

Choose a tag to compare

@tomaarsen tomaarsen released this 16 Jun 14:02

This minor version is a correctness- and robustness-focused release. It fixes a silent scoring bug for causal-LM rerankers, corrects several hard-negative mining and GIST loss edge cases, restores TSDAE on transformers v5, and adds Apple Silicon (MPS) support for the cached losses.

The headline fix affects chat-template models that read the final token position, i.e. causal-LM rerankers (like Qwen3-Reranker) and last-token-pooling embedders: when an over-long input was truncated, the chat template's trailing suffix (e.g. the assistant prefill the model scores from) was silently dropped, producing wrong scores with no error. There's also a forward-looking deprecation: loading local custom code without trust_remote_code=True now warns, and will require it from v6.0.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==5.6.0

# Inference only, use one of:
pip install sentence-transformers==5.6.0
pip install sentence-transformers[onnx-gpu]==5.6.0
pip install sentence-transformers[onnx]==5.6.0
pip install sentence-transformers[openvino]==5.6.0

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==5.6.0
pip install sentence-transformers[audio]==5.6.0
pip install sentence-transformers[video]==5.6.0

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==5.6.0

Fixed silently wrong scores when truncation drops chat-template suffixes (#3787)

Chat-template models render the full conversation to a flat string before tokenizing, so when the rendered input is longer than the tokenizer's model_max_length, the tokenizer truncates it from the right and drops the template's trailing suffix: the fixed tokens a template appends after the content, e.g. a prompt, instruction, [/INST], or a trailing EOS. For models that read the final token position, this silently corrupted the result:

  • causal-LM rerankers (e.g. Qwen/Qwen3-Reranker-0.6B) score a pair from the last token's yes/no logits, and
  • last-token-pooling embedders read the final hidden state.

When the suffix was truncated away, that final position landed mid-document instead of after the prefill, so the score or embedding came from the wrong place.

Transformer.preprocess now detects when truncation drops the suffix and splices it back onto the tail of each truncated row. Because the fix lives in the shared base Transformer, it applies across SentenceTransformer, CrossEncoder, and SparseEncoder. It's enabled by default and saved to the model configuration. Pass processing_kwargs={"chat_template": {"restore_suffix": False}} to opt back into raw truncation.

Hard-negative mining and GIST loss correctness (#3821, #3817, #3816)

A trio of correctness and scalability fixes for hard-negative mining and the GIST losses:

  • Sign-independent relative margin: mine_hard_negatives(relative_margin=...) and the margin_strategy="relative" branch of GISTEmbedLoss / CachedGISTEmbedLoss used a multiplicative threshold (positive * (1 - margin)) that only behaves correctly when the positive-pair similarity is positive. When that similarity was negative, the threshold moved the wrong way and let through false negatives: candidates more similar to the anchor than the true positive. The threshold is now positive - |positive| * margin, identical to before for positive scores but correct for negative ones.
  • Distributed positive masking in the GIST losses: with gather_across_devices=True and a non-zero margin, the false-negative suppression mask protected the wrong columns on ranks beyond the first (it ignored the per-rank offset into the gathered batch), which set the true positive's logit to -inf and produced a +inf loss. The mask now accounts for the cross-rank offset, so multi-GPU GIST training stays finite.
  • Memory-bounded mining without FAISS: mine_hard_negatives(use_faiss=False) (the default) materialized the full (queries × corpus) similarity matrix at once, which could OOM on large corpora. It now batches over the query axis (controlled by faiss_batch_size, default 16384), bounding peak memory while producing identical results.

TSDAE weight tying restored on transformers v5 (#3781)

transformers v5 removed the private PreTrainedModel._tie_encoder_decoder_weights helper that DenoisingAutoEncoderLoss (TSDAE) used to tie its separate encoder and decoder. As a stopgap, v5.5 raised a RuntimeError for the default tie_encoder_decoder=True on transformers >= 5.0.0, effectively breaking TSDAE there unless you pinned an older transformers or disabled tying. TSDAE now ships its own tying routine that shares storage between encoder and decoder, so it works on both transformers <5 and >=5 with the default settings.

Deprecation: loading local custom code without trust_remote_code (#3807)

Sentence Transformers has historically treated any local model directory as implicitly trusted: local custom code (e.g. modeling_*.py) loaded even with trust_remote_code=False, unlike transformers. This discrepancy might be unexpected, so loading local custom code this way now emits a FutureWarning, and from v6.0 it will require trust_remote_code=True like in transformers.

Apple Silicon (MPS) support (#3812, #3818)

Two fixes for training on Apple Silicon:

  • Cached losses on MPS: CachedMultipleNegativesRankingLoss and CachedGISTEmbedLoss crashed at construction on MPS because their RandContext used a CUDA-only RNG path. They now run on MPS with deterministic replay preserved.
  • Legacy fit path and SparseEncoder sparsity on MPS: the legacy model.fit(..., use_amp=True) path hard-coded CUDA's AMP GradScaler / autocast, and SparseEncoder sparsity statistics called to_sparse_csr(), which is unimplemented on MPS. Both now work on Apple Silicon.

Bug Fixes

  • Cast learning-to-rank loss logits to float32 in #3800: the listwise learning-to-rank losses scatter the model's logits into a float32 matrix, which crashed with a dtype mismatch when the model itself was in bf16/fp16 and trained without bf16=True/fp16=True (with those enabled, autocast outputs float32 logits, so the common path was unaffected). Logits are now upcast to float32 in the loss.
  • Don't override device_map placement with the device argument in #3823: loading with model_kwargs={"device_map": ...} previously placed the backbone via accelerate, then immediately moved it to the default device, defeating device_map. It now keeps the backbone in place (moving the other modules onto its device), and warns if both device and device_map are passed.
  • Collapse single-key multimodal dicts to the bare modality in #3779: a single-key input like {"image": img} was classified as a combined modality and rejected by models that support only that one modality (e.g. BGE-VL), with a self-contradicting error. It is now treated as the bare "image" input, which unblocks vision-retrieval benchmarks like MTEB that pass {"image": ...}.
  • Clarify unsupported-modality error messages in #3792: mixed or combined multimodal inputs on models that can't fuse modalities produced confusing errors (e.g. Modality 'message' is not supported). The errors are now scenario-specific and suggest what to do, such as encoding each modality separately.
  • Guard distributed APIs in get_device_name in #3798: on PyTorch builds where torch.distributed is present but unavailable (some ROCm and CPU-only builds), get_device_name() crashed with AttributeError: module 'torch.distributed' has no attribute 'is_initialized'. It now checks is_available() first, across all distributed call sites.
  • Fix OpenVINO static quantization for optimum-intel 2.0 / OpenVINO 2026 in #3814: export_static_quantized_openvino_model defaulted its calibration dataset to the bare id "glue", which the stricter Hub repo-id validation now rejects. The default is now the namespaced "nyu-mll/glue".

Examples, Documentation, and Notebooks

A batch of example and documentation modernization, mostly migrating example scripts off deprecated datasets script-loaders and bare ids onto maintained Hugging Face datasets so they run on datasets 4.x:

  • Fix example datasets that crash on datasets 4.x in #3782 (e.g. quora, nq_open, yahoo_answers_topics).
  • Migrate the MS MARCO examples to datasets in #3783.
  • Migrate the multilingual parallel-sentences data-prep scripts to datasets in #3784.
  • Migrate the Quora semantic-search and clustering examples to datasets in #3785.
  • Modernize the AugSBERT data-augmentation STSb scripts in #3806, also fixing the QQP cross-domain script to use the quora-duplicates labels it previously ignored.
  • Refresh the CLIP / image-search notebooks in #3780, and modernize the CLIP training notebook in #3805.
  • Fix a reStructuredText table in the ContrastiveTensionLoss docstring in #3788, and add a low-VRAM hardware note to the efficiency docs in #3802.
  • Fix package build warnings and errors in #3809, and a batch of Sphinx doc-build problems in #3811.

All Changes

  • [chore] Increment dev version by @tomaarsen in #3775
  • [fix] Collapse single-key multimodal dicts to bare modality by @tomaarsen in #3779
  • chore: enable Dependabot weekly GitHub Actions bumps by @hf-dependantbot-rollout[bot] in #3786
  • Bump the actions group with 3 updates by @dependabot[bot] in #3790
  • Fix example datasets that crash on datasets 4.x by @omkar-334 in #3782
  • CLIP Notebooks Refresh by @lbourdois in #3780
  • fix TSDAE ...
Read more

v5.5.1 - Small Multimodal patch

Choose a tag to compare

@tomaarsen tomaarsen released this 20 May 07:37

This patch release fixes a small quirk with multimodal inference when using single-key multimodal inputs like model.encode({"image": ...}).

Install this version with

# Training + Inference
pip install sentence-transformers[train]==5.5.1

# Inference only, use one of:
pip install sentence-transformers==5.5.1
pip install sentence-transformers[onnx-gpu]==5.5.1
pip install sentence-transformers[onnx]==5.5.1
pip install sentence-transformers[openvino]==5.5.1

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==5.5.1
pip install sentence-transformers[audio]==5.5.1
pip install sentence-transformers[video]==5.5.1

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==5.5.1

Bug fixed

Previously, inference like model.encode({"image": ...}) or model.encode([{"image": ...}, ...]) would be inferred as the ("image",) modality, which differed from the inferred modality of "image" for just model.encode(my_image) or model.encode([my_image, my_image_2, ...]).

This results in confusing errors if the model doesn't have a modality_config mapping for ("image",) in addition to "image", so now a single-key multimodal dict is collapsed to the bare modality (just "image" in this example).

This affected this code:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('BAAI/BGE-VL-base', trust_remote_code=True)
embedding = model.encode({"image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ettin-reranker/mteb_ndcg10_all-MiniLM-L6-v2.png"})
print(embedding.shape)

Which previously failed as the model only implements a path for "text", "image", and ("image", "text").

All Changes

  • [fix] Collapse single-key multimodal dicts to bare modality by @tomaarsen in #3779

Full Changelog: v5.5.0...v5.5.1

v5.5.0 - Training Agent Skill, EmbedDistillLoss, and ADRMSELoss

Choose a tag to compare

@tomaarsen tomaarsen released this 12 May 14:06

This release ships the train-sentence-transformers Agent Skill, adds two new training losses, and brings a long list of robustness and correctness fixes.

The new train-sentence-transformers Agent Skill lets AI coding agents (Claude Code, Codex, Cursor, Gemini CLI, ...) drive end-to-end training and fine-tuning across all three model types. EmbedDistillLoss is a new embedding-level knowledge distillation loss for SentenceTransformer: it aligns a student model's embeddings with pre-computed teacher embeddings, an alternative to the score-based distillation provided by MarginMSELoss and DistillKLDivLoss. ADRMSELoss is a new listwise learning-to-rank loss for CrossEncoder from the Rank-DistiLLM paper. encode() and predict() also gain a per-call processing_kwargs override, and more.

Install this version with

# Training + Inference
pip install sentence-transformers[train]==5.5.0

# Inference only, use one of:
pip install sentence-transformers==5.5.0
pip install sentence-transformers[onnx-gpu]==5.5.0
pip install sentence-transformers[onnx]==5.5.0
pip install sentence-transformers[openvino]==5.5.0

# Multimodal dependencies (optional):
pip install sentence-transformers[image]==5.5.0
pip install sentence-transformers[audio]==5.5.0
pip install sentence-transformers[video]==5.5.0

# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==5.5.0

The train-sentence-transformers Agent Skill (#3752)

If you use an AI coding agent (Claude Code, Codex, Cursor, Gemini CLI, OpenCode, ...), you can now install the train-sentence-transformers Agent Skill and ask your agent to fine-tune a model on your data:

hf skills add train-sentence-transformers              # installs under ./.agents/skills/
hf skills add train-sentence-transformers --global     # installs under ~/.agents/skills/
hf skills add train-sentence-transformers --claude     # also symlinks into .claude/skills/

The skill gives the agent curated, version-aware guidance for training SentenceTransformer (bi-encoder), CrossEncoder (reranker), and SparseEncoder/SPLADE models, covering base model selection, loss and evaluator choice, hard-negative mining, distillation, LoRA, Matryoshka, multilingual training, static embeddings, plus a set of production-ready training template scripts. Then you can prompt your agent with things like:

"Train a multilingual sentence-transformer on Dutch legal pairs."

"Fine-tune a cross-encoder reranker on (question, answer) pairs from my dataset, mine hard negatives, and push to my Hub repo."

"Train a German sparse embedding model with high sparsity."

"Can you train a static embedding model on 100k code triplets?"

The skill lives in the repository under skills/train-sentence-transformers/ and is mirrored to the huggingface/skills marketplace on each release.

New loss: EmbedDistillLoss (#3665)

Introduces EmbedDistillLoss (Kim et al., 2023), an embedding-level knowledge distillation loss for SentenceTransformer. Rather than distilling teacher scores (MarginMSELoss, DistillKLDivLoss), it directly aligns the student's sentence_embedding with a pre-computed teacher embedding passed via the dataset's label column. The comparison uses a configurable distance_metric, one of "cosine" (the default), "l2", or "mse". When the student and teacher dimensions differ, pass projection_dim=<teacher_dim> to add a learnable projection from the student's embedding space into the teacher's. That projection lives on the loss rather than on the saved model, so use loss.save_projection(...) / loss.load_projection(...) to reuse it across stages (e.g. like done in Arkam et al. for Jina v5). As part of this change, MSELoss is now a thin subclass of EmbedDistillLoss with distance_metric="mse", and also gains the optional projection_dim argument.

from datasets import Dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.sentence_transformer.losses import EmbedDistillLoss

student_model = SentenceTransformer("microsoft/mpnet-base")
teacher_model = SentenceTransformer("all-mpnet-base-v2")

train_dataset = Dataset.from_dict({
    "sentence": ["It's nice weather outside today.", "He drove to work."],
})

# Pre-compute teacher embeddings once and store them as the `label` column
def add_teacher_embeddings(batch):
    return {"label": teacher_model.encode(batch["sentence"]).tolist()}

train_dataset = train_dataset.map(add_teacher_embeddings, batched=True)

loss = EmbedDistillLoss(student_model, distance_metric="cosine")
# If the student and teacher dimensions differ, add a learnable projection:
# loss = EmbedDistillLoss(student_model, distance_metric="cosine", projection_dim=768)

trainer = SentenceTransformerTrainer(
    model=student_model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

See the updated model distillation examples and the loss overview for more.

New loss: ADRMSELoss for Cross Encoders (#3690)

Introduces ADRMSELoss (Approx Discounted Rank Mean Squared Error), a listwise learning-to-rank loss for CrossEncoder from the Rank-DistiLLM paper (Schlatt et al., ECIR 2025). It computes a differentiable approximation of each document's rank via pairwise sigmoids and minimizes the nDCG-discounted squared error against the true ranks derived from the labels. It expects listwise inputs: a (query, [doc1, ..., docN]) pair plus a [score1, ..., scoreN] label list per sample (binary or continuous labels, variable document counts allowed). It's designed for LLM-distillation reranking, where the per-document scores come from a strong LLM's ordering.

from datasets import Dataset
from sentence_transformers import CrossEncoder, CrossEncoderTrainer
from sentence_transformers.cross_encoder.losses import ADRMSELoss

model = CrossEncoder("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "query": ["What are pandas?", "What is the capital of France?"],
    "docs": [
        ["Pandas are a kind of bear.", "Pandas are kind of like fish."],
        ["The capital of France is Paris.", "Paris is the capital of France.", "Paris is quite large."],
    ],
    "scores": [[0.95, 0.1], [0.98, 0.92, 0.2]],
})
loss = ADRMSELoss(model)

trainer = CrossEncoderTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

There's a full MS MARCO example at training_ms_marco_adrmse.py. Note that LambdaLoss generally remains the strongest loss in the listwise family. See the Cross Encoder loss overview for guidance on picking a loss.

Per-call processing_kwargs override (#3753)

SentenceTransformer.encode() / encode_query() / encode_document(), SparseEncoder.encode(), CrossEncoder.predict(), and model.preprocess() now accept a processing_kwargs argument that overrides the processor/tokenizer kwargs configured at construction time, for a single call. It has the same nested structure as the processing_kwargs constructor argument (top-level keys text, audio, image, video, common, chat_template) and is shallow-merged on top of the instance-level settings, so you can override just one setting (e.g. max_length) and leave the rest intact.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# Override processor kwargs (e.g. max_length, truncation) for this call only:
embeddings = model.encode(
    ["a short text", "a much longer text that you want truncated more aggressively ..."],
    processing_kwargs={"text": {"max_length": 256, "truncation": True}},
)

This is especially handy for vision-language models, where you can change the image resolution per call, e.g. model.encode(images, processing_kwargs={"image": {"max_pixels": 256 * 256}}).

Smaller Features

  • Allow CrossEncoder module stacks that don't start with a Transformer, and recognize a trailing Dense(module_output_name="scores") as the scoring head, by @tomaarsen in #3742: num_labels now reads that head's out_features, and model.config / model.model return None when there's no underlying transformers model.
  • Infer that a model is an IR model on its generated model card when an InformationRetrievalEvaluator / NanoBEIREvaluator (or their sparse variants) was used during training, by @tomaarsen in #3741: the usage snippet then shows encode_query / encode_document, even without IR prompt names or a Router architecture.
  • Warn at model-load time when the installed transformers version is too old to honor use_bidirectional_attention / is_causal flags in a model's config (e.g. for google/embeddinggemma-300m), rather than silently ignoring them, by @tomaarsen in #3726.

Bug Fixes

  • Use the first non-pad token for CLS pooling with left-padding tokenizers by @tomaarsen in #3767: pooling_mode="cls" previously returned the embedding at position 0, which is a [PAD] token for left-padded inputs (common with decoder-only models), silently producing incorrect sentence embeddings. It now ...
Read more