pypi sentence-transformers 5.6.1
v5.6.1 - Flash Attention Fix for XLM-R and RoBERTa Models

5 hours ago

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

Don't miss a new sentence-transformers release

NewReleases is sending notifications on new releases.