Skip to content

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

Latest

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 processors, so e.g. processing_kwargs={"text": {"padding_side": "right"}} silently scored padding tokens. The produced attention mask itself is now verified, raising on the first batch that is actually padded on the wrong side.

Security policy and a v6.0 deprecation for third-party module imports (#3858, #3860)

The repository now has a Security Policy (#3858): please report vulnerabilities through GitHub's private vulnerability reporting, and see the policy for the threat model and scope.

In the same vein, module class imports are now trust-checked (#3860). A model's modules.json references its module classes by import path. Classes from sentence_transformers itself are always fine, but a reference into any other installed package (e.g. some_library.models.CustomModule) used to be imported without any trust gate. Loading such a model from a remote, untrusted source now emits a FutureWarning, and from v6.0 it will require trust_remote_code=True, in line with the v5.6.0 deprecation of local custom code (#3807). The WordEmbeddings module's configurable tokenizer_class now goes through the same gate. Additionally, passing a file path as model_name_or_path raises a clear NotADirectoryError up front instead of potentially loading a same-named Hub repository.

Evaluator fixes (#3886, #3878, #3853, #3883, #3861)

Several fixes for BinaryClassificationEvaluator and ReciprocalRankFusionEvaluator, which their sparse counterparts inherit:

  • Inverted Euclidean/Manhattan metrics (#3886): since v3.3.0, "euclidean" and "manhattan" in similarity_fn_names produced metrics for the opposite classifier: accuracy, F1, average precision, MCC, and the thresholds (which came out negative) were all wrong. Beyond reporting, with one of these as the first similarity function, the inverted average precision could steer metric_for_best_model/load_best_model_at_end during training. Expect euclidean/manhattan numbers to jump upward versus previous releases. Cosine and dot metrics were always correct.
  • Misaligned results CSV (#3878): the CSV data row silently skipped the *_accuracy_threshold and *_f1_threshold columns, shifting every subsequent value under the wrong header, including in the default cosine-only configuration. The returned metrics and logs were unaffected, only the CSV file was misaligned.
  • KeyError: 'cosine' (#3853): requesting multiple similarity functions without cosine (e.g. similarity_fn_names=["dot", "euclidean"]) crashed while aggregating the max_* metrics.
  • Canonical Reciprocal Rank Fusion (#3883): ReciprocalRankFusionEvaluator used 0-based ranks and gave every document a score contribution from both retrievers, even from a retriever that did not return it, under-ranking documents that both retrievers agree on. The fusion now matches canonical RRF, so fused rankings and metrics shift, generally upward (the documented hybrid search example went from 32.62 to 32.95 NDCG@10).
  • Unprefixed primary_metric (#3861): with a name set, ReciprocalRankFusionEvaluator prefixed its result keys but not primary_metric, so the standard results[evaluator.primary_metric] access raised KeyError, e.g. when used during training.

Hard-negative mining and semantic search fixes (#3871, #3872, #3873, #3904, #3887, #3888, #3907, #3909)

A batch of fixes for mine_hard_negatives and the semantic_search_* helpers:

  • Prompt-aware embedding cache (#3871): the cache_folder cache key ignored the prompt arguments, so mining runs with the same texts but a different query_prompt/corpus_prompt (or prompt names) silently reused stale embeddings. The prompts are now part of the key. As a one-time side effect, caches written by older versions are recomputed after upgrading. Resolves #3870.
  • Too few FAISS candidates, and -1 padding as candidates (#3872): the use_faiss=True path retrieved range_max + 1 candidates instead of range_max + max_positives, yielding fewer negatives than the default path on multi-positive datasets. Additionally, when the candidate window exceeded the corpus size, FAISS pads its results with index -1, which Python resolves to corpus[-1], so the last corpus document could be mined as a negative for any query. Padded slots are now disqualified.
  • Small-corpus crash without FAISS (#3873): the default path crashed with RuntimeError: selected index k out of range whenever range_max + max_positives exceeded the corpus size, easy to hit on modest corpora with default settings. It now mines as many negatives as the corpus allows, matching the FAISS path.
  • Up-front validation of the mining window (#3904): configurations where num_negatives cannot fit in the [range_min, range_max) window used to crash mid-run with opaque tensor shape errors, and a negative range_min did not crash at all: it silently sliced the candidate window from the wrong end and mined the easiest candidates as "hard" negatives. Both now raise a clear ValueError. The 2048-candidate retrieval cap now only applies to FAISS on GPU, where it also fixes a crash by accounting for max_positives, and FAISS on CPU is no longer capped. Resolves #3903.
  • Phantom results in semantic_search_faiss (#3887): with a corpus smaller than top_k, FAISS -1 padding leaked into the returned hits as {"corpus_id": -1} entries with garbage scores, could outrank real documents after rescoring, and segfaulted on an empty index with rescore=True. Result lists now only contain real documents, so they can hold fewer than top_k entries for small corpora.
  • Correct corpus_precision documentation (#3888): the semantic_search_faiss docs listed "int8"/"binary", but the function accepts "float32", "uint8", and "ubinary". The docs are corrected and unsupported values now raise an immediate ValueError.
  • Query alignment in semantic_search_seismic (#3907): a query that matched no documents crashed the result formatting with an IndexError. Results now stay aligned with the input query order, with an empty list for no-match queries. The corpus_index type hint and docstring now describe the bare SeismicIndex the function actually accepts (the documented tuple never worked).
  • Single queries in semantic_search_qdrant (#3909): a single 1D query embedding (as returned by encode_query for one text) was iterated as if each vocabulary entry were its own query, silently running one Qdrant search per vocabulary token and returning that many result lists instead of one. It is now treated as a batch of one, and query tensors that are neither 1D nor 2D raise a clear ValueError.

Multimodal input handling (#3876, #3841, #3897, #3866)

  • Per-sample video_metadata alignment (#3876): in a batch mixing videos with and without metadata, the batch-level metadata list came out shorter than the batch, so metadata silently attached to the wrong videos and, with frame sampling enabled, trailing videos could be dropped entirely. Metadata is now aligned per sample. Audio batches mixing conflicting sampling_rates also now raise instead of silently processing all audio at whichever rate came last. Resolves #3874.
  • 2-frame videos misrouted as text pairs (#3841): a video passed as a list of exactly two frame paths was routed into the query/document pair handling meant for text (one frame or three-plus frames worked fine). Only text inputs are treated as pairs now. Resolves #3840.
  • torchcodec import robustness (#3897): torchcodec raises RuntimeError (not just ImportError) when it is installed but unusable, e.g. on a broken FFmpeg setup, which made import sentence_transformers itself fail. It is now treated as an unavailable optional dependency, keeping text-only usage working. Resolves #3896.
  • Better error for sibling metadata keys (#3866): passing video_metadata or sampling_rate as a sibling key next to "video"/"audio" in a multimodal dict now produces an error that spells out the working nested form (e.g. {"audio": {"array": ..., "sampling_rate": 16000}}).

More training-loss correctness (#3827, #3868, #3838)

  • Exclude padding from the Plackett-Luce normalizer in ListMLELoss/PListMLELoss (#3827): padded list positions entered the normalizer with roughly unit mass each, so a query's loss and gradients depended on how much padding its batch happened to contain. Padded positions are now excluded before the normalizer. Retraining the documented MS MARCO recipes lifts NanoBEIR mean nDCG@10 from roughly 0.39 to 0.53 for ListMLELoss and from 0.514 to 0.525 for PListMLELoss.
  • Pairwise similarity default in SparseCoSENTLoss (#3868): the default similarity_fct was the matrix-valued util.cos_sim where CoSENT needs the pairwise similarity, and broadcasting kept the loss finite while silently optimizing a different objective. The default is now util.pairwise_cos_sim, matching the docstring, CoSENTLoss, and SparseAnglELoss. If you passed similarity_fct explicitly, you were unaffected.
  • Warn on non-binary labels in the contrastive losses (#3838): ContrastiveLoss and OnlineContrastiveLoss expect 0/1 labels. They now emit a one-time warning when given anything else, explaining the actual behavior: ContrastiveLoss uses such labels directly as term weights, and OnlineContrastiveLoss silently drops those pairs from the loss. Resolves #3382.

Bug Fixes

  • Fix optimum ONNX export in #3831: model.config became a read-only property in v5.5.0, which broke optimum's ONNX export (it assigns model.config while standardizing attributes). A setter now delegates to the underlying transformers model. Resolves #3830.
  • Fix an eval DataLoader worker leak with dataloader_persistent_workers=True in #3895: every evaluation built and prepared fresh eval DataLoaders whose persistent worker processes were never released, accumulating over long training runs until file descriptor exhaustion (Too many open files) or OOM. Prepared eval dataloaders are now cached and reused per eval dataset name, and evaluate()/get_eval_dataloader() accept a dataset name string.
  • Speed up community_detection and reduce its memory usage in #3832: intermediate communities are stored as compact uint32 arrays (about 5x less memory), GPU results are moved to CPU once per batch, and the overlap-removal step is vectorized. Outputs are identical, and dense-community workloads measured up to 3.8x faster.
  • Expand the community_detection candidate window on threshold ties in #3900: on the CPU path, members whose similarity equals the threshold exactly (e.g. duplicate detection with threshold=1.0 on identical or one-hot vectors) never triggered a window expansion, silently capping communities at an internal window size (often 50). Resolves #3899.
  • Don't mutate the input in select_max_active_dims in #3852: the utility zeroed non-top-k values in the caller's tensor in place. It now returns a fresh tensor, supports single 1D embeddings (previously a crash), and raises a ValueError for non-positive max_active_dims (previously returned all-zero embeddings for 0).
  • Support mixed sparse and dense inputs in the Euclidean and Manhattan similarities in #3906: model.similarity() and model.similarity_pairwise() with similarity_fn_name="euclidean" or "manhattan" crashed on a SparseEncoder when one side was encoded with convert_to_sparse_tensor=False (cosine and dot already handled the mix), as did pairwise_angle_sim. Mixed inputs now match the all-dense results.
  • Use the Sentence Transformer specific model card template in #3828: automatically generated SentenceTransformer model cards were built from the shared base template, missing the "Additional Resources" documentation links that CrossEncoder and SparseEncoder cards already had.
  • Declare tokenizers as a direct dependency in #3844: it is imported directly but was only pulled in transitively via transformers, which mattered for strict resolvers and minimal environments. Resolves #3519.

All Changes

  • [chore] Increment dev version by @tomaarsen in #3824
  • [fix] Pass axis=-1 to np.packbits in binary/ubinary quantization by @JSap0914 in #3825
  • [model_card] Use ST-specific model card template for ST by @tomaarsen in #3828
  • [fix] Exclude padding from the Plackett-Luce normalizer in ListMLE/PListMLE losses by @Incheonkirin in #3827
  • [tests] Fix tiny-random tests with newer transformers by @tomaarsen in #3829
  • Bump actions/checkout from 6.0.3 to 7.0.0 in the actions group by @dependabot[bot] in #3834
  • [ci] Exclude librosa/numba/llvmlite on Python 3.13 by @tomaarsen in #3835
  • Bump actions/setup-python from 6.2.0 to 6.3.0 in the actions group by @dependabot[bot] in #3846
  • [Fix] Add setter for config property to fix optimum ONNX export by @lcheng321 in #3831
  • fix: declare tokenizers as a direct dependency by @aadhar-build in #3844
  • [Fix] only call pair_to_messages for text pairs by @yushuosun in #3841
  • Warn when (Online)ContrastiveLoss gets non-binary labels by @oyinkanchekwas in #3838
  • Fix BinaryClassificationEvaluator KeyError('cosine') with multiple non-cosine metrics by @vineethsaivs in #3853
  • 🚨 Clip out-of-range values in int8/uint8 quantization by @anjaliy11 in #3847
  • [fix] Memory growth in community_detection with larger datasets by @AbdelRahmanYaghi in #3832
  • [tests] Rename a reranker test model from v6 to v54 by @tomaarsen in #3857
  • add tests for append_to_last_row by @RavSinghChandan in #3855
  • Fix select_max_active_dims mutating its input tensor in place and correct its docstring by @teddytennant in #3852
  • [Fix] Make model.compile() apply to encode/predict by @tomaarsen in #3848
  • [security] Create a Security Policy by @tomaarsen in #3858
  • [security] Trust-check every non-sentence-transformers module class import by @tomaarsen in #3860
  • Fix ReciprocalRankFusionEvaluator primary_metric not being prefixed by @vineethsaivs in #3861
  • [tests] Skip bf16 + Windows + CPU forwards, as they can WindowsError on torch 2.13 by @tomaarsen in #3863
  • [losses] Consolidate GradCache into one shared engine, fix silently wrong gradients, add mini_batch_num_tokens by @tomaarsen in #3862
  • Improve the error message when video_metadata is passed as a sibling modality key by @tomaarsen in #3866
  • Bump astral-sh/setup-uv from 8.2.0 to 8.3.2 in the actions group by @dependabot[bot] in #3867
  • [fix] Use a pairwise similarity as the SparseCoSENTLoss default by @vineethsaivs in #3868
  • [warn] Check causal left padding against the attention mask by @tomaarsen in #3869
  • Include query/corpus prompts in mine_hard_negatives cache key by @ErenAta16 in #3871
  • Fix FAISS hard-negative mining retrieving too few candidates by @Osamaali313 in #3872
  • fix: prevent mine_hard_negatives crash when candidates exceed corpus size by @Kropiunig in #3873
  • 🚨 Fix position_ids offsetting for RoBERTa-family models when flattening inputs by @tomaarsen in #3879
  • Fix misaligned results CSV in BinaryClassificationEvaluator by @vineethsaivs in #3878
  • Bump the actions group with 3 updates by @dependabot[bot] in #3891
  • docs: add missing docstring to to_scipy_coo by @RavSinghChandan in #3843
  • Keep per-sample video_metadata aligned with the batch, raise on conflicting audio sampling rates by @meutsabdahal in #3876
  • Fix torchcodec runtime import failures on partial install by @mturac in #3897
  • [fix] RRF scoring: only rank retrievers that returned a document, and use 1-based ranks by @eSVeeF in #3883
  • tests: mock Hub model and dataset existence checks to prevent rate-limit CI failures by @tomaarsen in #3898
  • [fix] Detach KL teacher in AdaptiveLayerLoss by @Kaif10 in #3880
  • 🚨 Add layer_weighting option to AdaptiveLayerLoss and Matryoshka2dLoss by @tomaarsen in #3901
  • [ce] Raise when a chat template cannot carry query/document pairs by @tomaarsen in #3902
  • Fix inverted Euclidean/Manhattan metrics in BinaryClassificationEvaluator by @Kropiunig in #3886
  • Drop FAISS padding placeholders from semantic_search_faiss results by @ErenAta16 in #3887
  • Correct the corpus_precision values documented for semantic_search_faiss by @ErenAta16 in #3888
  • Expand the community window on ties with the threshold by @LK-maker-007 in #3900
  • [trainer] Fix eval DataLoader worker leak with dataloader_persistent_workers by @mjun0812 in #3895
  • [fix] Reject num_negatives larger than the range_min/range_max window in mine_hard_negatives by @LK-maker-007 in #3904
  • [fix] Support mixed sparse/dense inputs in euclidean and manhattan similarity by @LuShadowX in #3906
  • Keep queries aligned in semantic_search_seismic when a query matches nothing by @LK-maker-007 in #3907
  • [fix] Treat a 1-dimensional query embedding as a single query in semantic_search_qdrant by @LK-maker-007 in #3909
  • [fix] Name the lone suggested parameter in the mine_hard_negatives missing-negatives message by @tomaarsen in #3912

New Contributors

Special Thanks

Beyond the pull request authors listed above:

Full Changelog: v5.6.1...v5.7.0