Matryoshka

[yamiteru]

Are there any plans on supporting Matryoshka embeddings?

[stephantul]

Hello!

We actually support this already, you can pass in dimensionality when loading the model, like so:

from model2vec import StaticModel

s = StaticModel.from_pretrained("minishlab/potion-base-8m", dimensionality=32)

Any distilled model that was distilled with PCA will support dimensionality reduction like this out of the box. Let me know if this helps.

[yamiteru]

I have this code:

from model2vec import StaticModel

text = "I really like apples!"

potion32_32 = StaticModel.from_pretrained("minishlab/potion-base-32m", dimensionality=32)
potion32_64 = StaticModel.from_pretrained("minishlab/potion-base-32m", dimensionality=64)

print(potion32_32.encode(text))
print(potion32_64.encode(text))

It prints this:

[-0.2889684  -0.45255008 -0.25858387 -0.09204515 -0.34086925  0.31868082
  0.21026841  0.18883598  0.24139526  0.04028355 -0.14286123 -0.05872613
  0.21438909 -0.00420785 -0.09024969 -0.18836147 -0.00819852 -0.14227319
  0.08295885 -0.0453723  -0.00551837 -0.02566796  0.15842609  0.13274904
  0.09205312 -0.03018598 -0.03111458  0.06423686  0.0348335  -0.26033974
 -0.03807567  0.08831255]
[-0.20676827 -0.3238174  -0.18502694 -0.06586193 -0.24390537  0.22802867
  0.15045533  0.13511957  0.17272782  0.02882447 -0.10222283 -0.04202086
  0.15340383 -0.00301088 -0.06457721 -0.13478005 -0.00586636 -0.10180207
  0.05936033 -0.03246566 -0.00394861 -0.01836644  0.11336011  0.09498716
  0.06586763 -0.02159925 -0.0222637   0.045964    0.02492474 -0.18628334
 -0.02724464  0.06319111  0.07593099  0.04000881 -0.19117253  0.05237667
 -0.16452202  0.10882624  0.0757569  -0.06416594  0.04851569 -0.0323696
  0.03383688  0.20155692  0.19408153  0.04402791  0.12130599  0.2515158
  0.12770426 -0.02624993  0.03664309  0.19524474 -0.05818356 -0.02506019
 -0.13721164  0.19186205 -0.06991481 -0.08733629  0.15303712 -0.11904417
  0.06560926 -0.0747398  -0.13996041  0.18473662]

As you can see the first 32 numbers in these two vectors are not equal.

This makes it impossible to query top-k based on the rough meaning of vectors and later do a more fine-grained query on top of the results.

[yamiteru]

If my understanding is correct I might be able to get Matryoshka our of model2vec if I try to distill a model that already outputs Matryoshka embeddings.

[stephantul]

This is because normalize is set to True. If you set it to False, you will get the same embeddings.

[yamiteru]

Oh I see! Have you observed score degradation in let's say MTEB using normalize=False?

[stephantul]

Yes, you'd get a bit lower accuracy, unfortunately. So what I would do, assuming you want to do something like retrieval, is the following:

1. set normalize to False, run inference once, this gives you a vector v
2. store v, and the v[:32] (or whatever dimension you desire), and normalize both of them separately
3. this will give you fast inference time and good performance but at a slightly higher storage overhead

Alternatively:

2. store v and compute the norm for v and v[:32]
3. when running inference, you can still compute the dot product as cosine using the precomputed norms.

This will be slower, but will not store vectors twice.

This aside, if you want to store many vectors and want to use HNSW or something, you will need to create 2 separate indices anyway, because you can't an "slice" HNSW index.