Compatibility with MUVERA encoding

[davidmezzetti]

First off, excellent work with this project and all the hard work you guys have put into late interaction models!

I've been working to add MUVERA to txtai and I've been having difficulty getting any PyLate models to work well with MUVERA encoding.

I also asked this question over on the muvfde repo and used those benchmark scripts to rule out anything I'm doing wrong (results copied here for clarity).

Metric	ModernColBERT	FDE(20,5,16) ModernColBERT
ndcg_at_1	0.35294	0.05882
ndcg_at_3	0.31416	0.05294
ndcg_at_5	0.28701	0.05297
ndcg_at_10	0.26125	0.05245
ndcg_at_20	0.24186	0.05263
ndcg_at_100	0.247	0.0671
ndcg_at_1000	0.33668	0.16921
map_at_1	0.03499	0.00202
map_at_3	0.06291	0.00427
map_at_5	0.0735	0.00554
map_at_10	0.08719	0.00736
map_at_20	0.09664	0.00914
map_at_100	0.11057	0.01303
map_at_1000	0.12253	0.02048
recall_at_1	0.03499	0.00202
recall_at_3	0.07216	0.00647
recall_at_5	0.09391	0.01138
recall_at_10	0.1267	0.01995
recall_at_20	0.15834	0.03348
recall_at_100	0.26555	0.09853
recall_at_1000	0.57746	0.4513

Curious if anyone has tried MUVERA encoding and has insights into parameters that work. Or if there is something fundamentally different with these embeddings vs the ones generated with models trained with the Stanford code.

Thank you!

[NohTow]

Hey,
Thanks for the kind words!
Those results are indeed kind of bad!

I am not that well versed with the actual internals of MUVERA but there is actually no reason why PyLate models specifically would be bad with it. One thing you could try is to run the benches with PyLate but using the models from the already reported benches (we support loading any stanford models), the results should match, if not, then it is most likely just an issue with the connexion between the lib and PyLate.
I'll ping the author of the lib as well (which is a wrapper of the Google wheels), he might know what's wrong!

[davidmezzetti]

Thank you for getting back. I did locally run the ColBERT model in PyLate format and the scores were close to the Stanford weights. I also tried JinaColBERT and that did well too.

This just made me wonder if this algorithm really only works well with the Stanford models. I'd ask the same question to the authors of MUVERA but there doesn't seem to be an easy way to do that.

[NohTow]

Uh, there is 0 reason I could think of that would make the models trained with PyLate worse for MUVERA: the objective is the exact same!
Could you quickly try ColBERT-small? This is the exact same training setup as GTE-ModernColBERT, the only difference is the base model.

I think I have a guess though, we agree that the results you reported for GTE-ModernColBERT without MUVERA is the one we report and not something you reproduced?
If yes, I suspect that you are running those tests on CPU but GTE-ModernColBERT inference is broken on CPU, see #125.
I really need to fix this (I have a branch working, but I want to make sure I do not miss anything for training).

[NohTow]

If you want a monkey patch for inference on CPU with GTE-ModernColBERT, what you can do is changing the line

# We keep all tokens in the query (no skiplist) and we do not want to prune expansion tokens in queries even if we do not attend to them in attention layers
masks = torch.ones_like(
input=out_features["input_ids"], dtype=torch.bool
)

By

# We only keep the original tokens and prune padding tokens
masks = out_features["attention_mask"].bool()

The only issue is that then the queries embeddings are not padded to the same length, so you need to do this to get the query_embeddings that you can feed to the index for search:

 query_embeddings = torch.nn.utils.rnn.pad_sequence(
                model.encode(
                    self.queries,
                    prompt_name=self.query_prompt_name,
                    prompt=self.query_prompt,
                    batch_size=self.batch_size,
                    is_query=True,
                    show_progress_bar=self.show_progress_bar,
                    convert_to_tensor=True,
                ),
                batch_first=True,
                padding_value=0,
            )

I am deeply sorry, this should have been fixed a long time ago but I've been lacking time and I am afraid of breaking things without more tests

[davidmezzetti]

I tried the same script on ColBERT-small and got this:

Metric	ColBERT-small	FDE(20,5,16) ColBERT-small
ndcg_at_1	0.32817	0.04334
ndcg_at_3	0.29374	0.04055
ndcg_at_5	0.28046	0.04352
ndcg_at_10	0.25628	0.04326
ndcg_at_20	0.2416	0.04326
ndcg_at_100	0.24336	0.05889
ndcg_at_1000	0.3362	0.15342
map_at_1	0.03508	0.0022
map_at_3	0.05332	0.00405
map_at_5	0.06508	0.00561
map_at_10	0.07749	0.00739
map_at_20	0.08749	0.00863
map_at_100	0.10264	0.0114
map_at_1000	0.11564	0.018
recall_at_1	0.03508	0.0022
recall_at_3	0.06317	0.00574
recall_at_5	0.08905	0.01115
recall_at_10	0.11913	0.02152
recall_at_20	0.15995	0.03204
recall_at_100	0.27193	0.08405
recall_at_1000	0.6013	0.39361

Just to be clear, I'm not suggesting anything is wrong with PyLate, I'm more questioning how generalizable the MUVERA algorithm is to models besides the standard ColBERT model.

[NohTow]

This is very weird.
IIUC:

• Stanford models work correctly with MUVERA, both with PyLate and stanford-nlp. -- These ones are trained using stanford
• Jina model works correctly with MUVERA, both with PyLate and stanford-nlp. -- This one is trained using stanford
• ColBERT-small does not work with PyLate (does it with stanford?). -- This one is trained using stanford
• GTE-ModernColBERT does not work with PyLate (we cannot test with stanford here) -- This one is trained using PyLate

What's very odd is that ColBERT-small is not working while it is trained with stanford-nlp, so it should work as well. Maybe this one is because the dimension of the embeddings is different (and too small), but I cannot figure out why GTE would not be working properly then. Are you running it on CPU?

[davidmezzetti]

This is very weird. IIUC:

* Stanford models work correctly with MUVERA, both with PyLate and stanford-nlp. -- These ones are trained using stanford

* Jina model works correctly with MUVERA, both with PyLate and stanford-nlp. -- This one is trained using stanford

* ColBERT-small does not work with PyLate (does it with stanford?). -- This one is trained using stanford

* GTE-ModernColBERT does not work with PyLate (we cannot test with stanford here) -- This one is trained using PyLate

What's very odd is that ColBERT-small is not working while it is trained with stanford-nlp, so it should work as well. Maybe this one is because the dimension of the embeddings is different (and too small), but I cannot figure out why GTE would not be working properly then. Are you running it on CPU?

I wouldn't focus on any of the standard multi-vector scores. These comments solely apply to MUVERA.

In summary, MUVERA seems to work with the standard model weights provided by Stanford and Jina's ColBERT model. Doesn't seem like any other models work well. Ultimately, I was curious if you had any insights from a PyLate standpoint.

It could just be a hyperparameter/tuning thing. But I've tried a lot of options for the MUVERA params and nothing really breaks the NDCG@10 threshold of 0.05 for any model besides the two mentioned above.

[NohTow]

It's very very strange, because it does not seems to be related to PyLate inference (because you manage to get good scores with some of the models) nor to be from the PyLate training (because ColBERT-small is bad as well).
Also, AFAIK all of these models are trained with KD using cross-encoders, so even there I am not sure why a particular model would be better or worse in this regard.
I know @bclavie did some tests on MUVERA (and he is the author of ColBERT-small), so maybe he has some insights as well?

[davidmezzetti]

More to the mystery here.

I trained a PyLate model using the standard contrastive learning example. It performs much better with MUVERA even though this model is designed for medical data vs being a generic model.

Metric	FDE(20,5,16) PubMedBERT ColBERT
ndcg_at_1	0.30805
ndcg_at_3	0.28178
ndcg_at_5	0.26326
ndcg_at_10	0.23928
ndcg_at_20	0.21769
ndcg_at_100	0.21301
ndcg_at_1000	0.30702
map_at_1	0.03329
map_at_3	0.05616
map_at_5	0.06669
map_at_10	0.07728
map_at_20	0.0839
map_at_100	0.09398
map_at_1000	0.10515
recall_at_1	0.03329
recall_at_3	0.06608
recall_at_5	0.08951
recall_at_10	0.12092
recall_at_20	0.14343
recall_at_100	0.22799
recall_at_1000	0.54349

Is it possible that contrastive vs the knowledge distillation objective is a factor?

[NohTow]

I was wondering at some point, but AFAIK Jina-ColBERT-v2 is also further fine-tuned using KD to distill a teacher...
So it does not seems like it is the reason

[davidmezzetti]

Interesting.

If there is anything I can do to help, please let me know. Also please feel free to close if this seems less conclusively related to PyLate.