github huggingface/sentence-transformers v6.3.0
v6.3.0 - Token-pooling-aware training and `kernels` support

one hour ago

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 SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="neuron")
embeddings = model.encode(["Native inference on AWS Neuron."])

For multi-process encoding, load the model on CPU in a fresh parent process so workers can use separate Neuron cores. See the PR for setup and device-allocation details.

Smaller Features

  • More efficient community detection: a shared batched implementation reduces GPU memory use on large embedding collections and speeds up CPU clustering. (#4342)
  • Faster utilities on Apple Silicon: community detection and hard-negative mining spend less time synchronizing with the device, improving runtime on MPS. (#4265, #4290)
  • Faster Hugging Face Dataset inputs: batch conversion of Dataset columns reduces data preparation overhead in encoding, prediction, mining, and selected sampler and evaluator workflows. (#4347)
  • Clearer model merging diagnostics: configuration errors explain incompatible inputs more clearly, while new benchmark examples help you compare merged checkpoints. (#4315)

Compatibility Notes

  • Multi-vector loss scoring: the four built-in contrastive and distillation losses now follow the model's similarity even with kernels disabled, so MeanMaxSim models use mean scores instead of the previous MaxSim sums. Check your contrastive scale and teacher-score convention when upgrading. To retain the old MaxSim behavior, set similarity_fct to colbert_scores for ranking losses, colbert_scores_pairwise for margin MSE, or colbert_kd_scores for KL distillation. (#4344)
  • Evaluation scores: historical scores can change because reranking and paraphrase-mining now follow the model's similarity, while CrossEncoder and RRF normalize MAP/nDCG against all known positives. Unretrieved positives receive no credit, so scores reflect retrieval failures. (#4252, #4268, #4313, #4054, #4114)
  • Checkpointing names: use activation_checkpointing_enable, activation_checkpointing, and activation_checkpointing_kwargs across supported Transformers versions, or keep the legacy gradient-checkpointing names in existing scripts. (#4328)
  • Saved compression settings: a constructor truncate_dim is saved with the model, but token_pooling is a runtime default, so supply the pooler again after reload. (#4286)
  • Pooler normalization: newly constructed hierarchical poolers normalize their output vectors by default. Older saved configurations without normalize_embeddings retain their unnormalized behavior to preserve compatibility with existing embeddings. (#4286)

Bug Fixes

Evaluation and ranking metrics

  • Use the model's configured similarity in reranking and paraphrase-mining evaluation, fixing dot-product and sparse-model scoring. (#4252, #4268, #4313)
  • Correct CrossEncoder and RRF MAP/nDCG normalization when retrieval misses positives, and handle empty CrossEncoder retrieval results. (#4054, #4114, #4264)

Encoding and low-precision models

  • Restore low-precision workflows by casting Router modules to their route's dtype and aligning distillation inputs and targets with the projection-weight dtype. (#4314, #4348)
  • Make CSR sparse autoencoders picklable so they work with multi-process encoding and DataLoader workers. (#4306)
  • Prevent StaticEmbedding from sharing and modifying its source tokenizer's backend. (#4253)
  • Accept pathlib.Path model paths in SentenceTransformer. (#4326)

Training

  • Preserve positional CrossEncoder constructor arguments after device. (#4350)
  • Stack tensor-formatted labels in the CrossEncoder data collator. (#4257)
  • Include evaluation loss components in logs when using custom metric prefixes. (#4237)
  • Exclude self-pairs when GIST guide scores tie. (#4284)

Hard-negative mining and retrieval examples

  • Keep anchors out of their own mined negatives and apply CrossEncoder rescoring even when no score filters are configured. (#4267, #4262)
  • Handle mining inputs and edge cases that previously failed, including a separate corpus supplied as a Dataset column and filters that exclude every candidate. (#4318, #4292)
  • Honor rescoring when reusing float USearch indexes. (#4251)

All Changes

  • [chore] Increment the dev version by @tomaarsen in #4245
  • Allow underscores in the NanoBEIREvaluator aggregate_key by @LindseyZ1205 in #4250
  • Restore the MLM training example on Transformers v5 by @Bembaby in #4247
  • Fix rescoring gate for reused USearch float indexes by @Aurantium-47 in #4251
  • [fix] Score SparseRerankingEvaluator with the model's similarity function by @mottopanikeiku in #4252
  • Bump the actions group with 2 updates by @dependabot[bot] in #4255
  • [fix] Stack tensor-formatted labels in the CrossEncoder data collator by @MaxFreedomPollard in #4257
  • Correct the DDP evaluator note in the training overviews by @Bembaby in #4248
  • Raise an error in fit() when InputExamples have different numbers of texts by @zikrulihsan in #4239
  • [examples] Fix SentenceTransformerTrainingArguments import in paraphrase and multi-task examples by @astroknot-sheep in #4246
  • Don't share the transformers tokenizer's backend in StaticEmbedding by @LindseyZ1205 in #4253
  • Honor CrossEncoder rescoring without mining score filters by @Aurantium-47 in #4262
  • Speed up community_detection ~64x on MPS by using the batched path by @m4zar in #4265
  • [tests] Simplify CrossEncoder mining regression tests by @tomaarsen in #4277
  • [fix] flaky rendezvous in distributed tests by @eSVeeF in #4266
  • [fix] Score RerankingEvaluator with the model's similarity function by @patrasq in #4268
  • [fix] Skip the anchor itself when mining hard negatives by @mottopanikeiku in #4267
  • 🚨 [fix] Strip left padding from token embeddings in encode() by @alejandrorodriguezalvarez884-dot in #4279
  • Log eval loss components for any metric_key_prefix by @LindseyZ1205 in #4237
  • Speed up minimum version tests by removing unnecessary FFmpeg installation by @tomaarsen in #4283
  • [fix] Pin documentation source links to the build commit by @j341nono in #4280
  • Handle empty retrieval results in cross-encoder reranking evaluation by @Aurantium-47 in #4264
  • [fix] Restore CrossEncoder.fit() with a custom loss_fct by @Kropiunig in #4285
  • Document adaptive-layer evaluation during training by @Bembaby in #4261
  • Exclude self-pairs when GIST guide scores tie by @Aurantium-47 in #4284
  • fix: preserve caller options during CrossEncoder initialization by @WitBlue6 in #4287
  • [docs] Fix quantization and SentenceTransformer API reference links by @pratikgx in #4291
  • Fix division by zero when mining filters exclude all negatives by @iBreaker in #4292
  • [fix] Warn when saving a Dense or CrossEncoder activation that is a plain function by @tanishqbaweja in #4289
  • Speed up mine_hard_negatives up to 10x on MPS by vectorizing the per-query loops by @sebastianbreguel in #4290
  • Add token pooling-aware and Matryoshka training for MultiVectorEncoder by @tomaarsen in #4286
  • Bump the actions group with 2 updates by @dependabot[bot] in #4295
  • [docs] Fix dead dataset, modules.json and reference links by @sarankumar1325 in #4300
  • [fix] Fix CrossEncoder.fit() crash with a DataLoader batch_sampler by @Kropiunig in #4296
  • [fix] Save the best model and call callback after each epoch in CrossEncoder.fit() by @tanishqbaweja in #4294
  • fix: preserve padded sequence lengths in LSTM outputs by @Afloat16 in #4301
  • [fix] Wait for Qdrant upload before searching by @Goodnight77 in #4305
  • Cap reranked scores by the positives the first stage missed by @vineethsaivs in #4054
  • [fix] Make SparseAutoEncoder picklable so CSR models work with multi-process encoding and DataLoader workers by @tanishqbaweja in #4306
  • Exclude unretrieved positives from observed RRF ranking metrics by @Afloat16 in #4114
  • [core] Add native AWS Neuron inference support by @JingyaHuang in #4293
  • Add model merging benchmarks and configuration diagnostics by @tomaarsen in #4315
  • [fix] Accept pathlib.Path as model_name_or_path in SentenceTransformer by @yuvraj-del in #4326
  • Exclude padded token embeddings before CNN convolution by @Afloat16 in #4323
  • [fix] Support evaluators that return a dictionary of metrics in SentenceTransformer.old_fit() by @ramiscaleai in #4320
  • Fix mine_hard_negatives crash on a separate corpus passed as a datasets v5 Column by @vineoy in #4318
  • Support activation checkpointing names with legacy aliases by @tomaarsen in #4328
  • Prefer MIC checkpoints and expand transformer architecture coverage by @tomaarsen in #4329
  • [fix] Cast the modules in each Router route to the route's dtype so Router models work in float16 and bfloat16 again by @tanishqbaweja in #4314
  • Bump astral-sh/setup-uv from 10.3.2 to 10.3.3 in the actions group by @dependabot[bot] in #4330
  • Add KURE-Sparse and MILCO to the SparseEncoder pretrained models by @tomaarsen in #4333
  • docs: fix grammar in semantic search example by @ddlw in #4340
  • fix: broadcast shared masks before pairwise MaxSim chunking by @flysheep-ai in #4332
  • [fix] Score ParaphraseMiningEvaluator with the model's similarity function by @MaiAditya in #4313
  • Ensure non-negative distances and zero diagonal in cosine_distance by @davrens101 in #4334
  • Fix CrossEncoder dropping positional args after device by @Ayush701-code-zm in #4350
  • Align distillation projection inputs and targets with projection-weight dtype for low-precision models by @Kropiunig in #4348
  • docs: Add OpenSearch asym-v1 to pretrained sparse models by @tomaarsen in #4351
  • Speed up dataset column to list conversion throughout ST by @deadstrobe5 in #4347
  • Remove commented Transformers usage sections from model card templates by @tomaarsen in #4352
  • Add kernels support for MaxSim, and Ward clustering, and transformers backbones by @tomaarsen in #4344
  • Reduce community detection memory and share batched top-k search across devices by @f4roukb in #4342
  • [docs] Update homepage hint for v6.3 by @tomaarsen in #4357

New Contributors

Full Changelog: v6.2.0...v6.3.0

Don't miss a new sentence-transformers release

NewReleases is sending notifications on new releases.