Integrate `SparseEncoder` model from SentenceTransformers `v5`

[arthurbr11]

Hello!

Context

We released Sentence Transformers v5 yesterday, introducing a new SparseEncoder model category. These are a subclass of SentenceTransformer that output sparse vectors instead of dense embeddings.

Feature Description

Add support for SparseEncoder models in MTEB. I tested the current version and found one main blocking issue in the Sentence TransformerWrapper to make it compatible :

mteb/mteb/models/sentence_transformer_wrapper.py

Line 114 in f27648b

embeddings = embeddings.cpu().detach().float().numpy()

with the .numpy (applied on a sparse tensor it's breaking and not wanted anyway)

Questions

1. Is there a specific reason for the .numpy() conversion? Can we make it conditional for sparse models?
2. For similarity computation, sparse models have an optimized similarity function that should be use this instead of MTEB's standard similarity functions, is it gonna be always the case ? Because if dense way of computing the similarity is used it will broke or be really long.
3. Since no sparse indexes exist in MTEB, encoding by chunk would be the best option to make sure the similarity compute is reasonable but is it plan to handle other indexes ?

This would probably be part of a bigger refactor of the Sentence Transformers handling with #2871 and other possible feature.

Would love to hear your thoughts thist.

cc @tomaarsen

Arthur Bresnu

[Samoed]

1. What embedding type SparseEncoder expecting?
2. We've changed this in v2 branch to use models similarity function if available, so I think this is resolved in v2

   mteb/mteb/models/abs_encoder.py

   Lines 35 to 53 in ed69e60

   	def similarity(self, embeddings1: Array, embeddings2: Array) -> Array:
   	if self.mteb_model_meta is None or (
   	self.mteb_model_meta is not None
   	and self.mteb_model_meta.similarity_fn_name is None
   	):
   	if (
   	hasattr(self, "model")
   	and hasattr(self.model, "similarity")
   	and callable(self.model.similarity)
   	):
   	return self.model.similarity(embeddings1, embeddings2)
   	return cos_sim(embeddings1, embeddings2)
   	if self.mteb_model_meta.similarity_fn_name is ScoringFunction.COSINE:
   	return cos_sim(embeddings1, embeddings2)
   	elif self.mteb_model_meta.similarity_fn_name is ScoringFunction.DOT_PRODUCT:
   	return dot_score(embeddings1, embeddings2)
   	elif self.mteb_model_meta.similarity_fn_name is ScoringFunction.MAX_SIM:
   	return max_sim(embeddings1, embeddings2)
   	raise ValueError("Similarity function not specified.")

3. Yes, we have plans for new indexes Add Support for Index-Based Retrieval for ColBERT #1593 Create interface with search #2728

[arthurbr11]

1. SparseEncoder models output a torch.sparse_coo by default.

2. That's great!

3. Nice! That means having one optimized for sparse vectors would be beneficial, it would be faster than the current model similarity computation.

[KennethEnevoldsen]

@arthurbr11 and @Samoed added a mock implementation of support for sparse embeddings (#2878). It seems like it will require a lot of adaptations, as the sparse torch matrix does not live up to Python's array API spec (see related issue here: #941).

@arthurbr11 would a solution be to convert the matrix to dense for cases such as classification, clustering etc., but keep it as a sparse vector for cases involving similarity measures?

I might be missing some important aspects (or smart tricks, so do let me know if there is anything I missed).

[arthurbr11]

You're right that we face compatibility challenges with sparse PyTorch matrices not conforming to Python's array API specification. To work around this, particularly for slicing operations, we've structured our encoder to return a list of sparse tensors instead of a single 2D sparse tensor when we want to do these operations. You can see our usage of this here and here the reason why we did it .

Converting to dense would indeed be the simplest solution, but it's not practical at all. The memory would explode really quickly. We're working with tensors that can exceed 30K of dim with only hundreds of non zeros values.

For the task you mentioned, i'm not sure i really understand why you can't use the native sparse adapt similarity of the model and then do operation on the scores ? But i'm not really aware of how these tasks evaluation work and i'm probably missing some difficulties.

[KennethEnevoldsen]

Ahh, that's great to hear - I think that could potentially solve the issues. We would have to expand our notion of an array a bit, but I think it might work.

Converting to dense would indeed be the simplest solution, but it's not practical at all. The memory would explode really quickly. We're working with tensors that can exceed 30K of dim with only hundreds of non zeros values.

For the task you mentioned, i'm not sure i really understand why you can't use the native sparse adapt similarity of the model and then do operation on the scores ? But i'm not really aware of how these tasks evaluation work and i'm probably missing some difficulties.

Ahh the classification and clustering etc, rely on the sklearn interface which requires a dense vector (or a sparse numpy vector)

[arthurbr11]

I think you were refering to sparse scipy vectors as there isn't numpy ones. If they are handled, then I think we can easily have this format from the torch sparse ones.

[KennethEnevoldsen]

Perfect, that seems to be the best way then - this is currently not first on my to-do list though (however, it is good to have an overview of what the problem entails). First priority is v2, and then we can add this in a backward-compatible manner.