Dynamic embedding size
Hello,
the model card says that this model can be configured based on 384, 512, 768, 1024, or 2048.
Could you please specify additional details, i.e. how to do it correctly? By modifying the average_pool function in the example code? By another way (e.g. by normalization of cropped embeddings)?
Thank you in advance!
Pavel
Hi Pavel,
I believe the model was trained with Matryoshka Representation Learning in the regime where we slice the embedding vector during training to predefined sizes, calculate the score/loss for each dimension, and subsequently sum them, meaning no projection layers were used.
In order to use the model at lower embedding sizes, you just need to slice the resultant embedding to the dimensionality you are after, with a batch of embeddings that should be embs[:, :desired_size].
The model was evaluated at the specified embedding dimension sizes, however the beauty of MRL is that it should also work well at all sizes from 384 to 2048!
You are however correct that the cropping should take place before normalization, so possibly in this step:
embedding = F.normalize(embedding[:,:desired_size], dim=-1)
Depending on how aggressively you crop the embeddings, you will see some performance decrease, but it should be relatively small. Could you please let us know if that is what you are seeing on your data?
Best,
Radek
Hi Radek,
thanks for your explanation.
The important detail is that mentioned sizes (384, 512, 768, 1024, or 2048) were used only for training while the inference code (your cropping+normalizing proposal) is not present in the model card.
This confused me initially -- but meanwhile, I tested it and now can confirm that desired_size=512 works well after normalization.
Thanks again. Best,
Pavel