Integrate with Sentence Transformers via MultiVectorEncoder

#2
by tomaarsen HF Staff - opened

Hello!

As mentioned in https://huggingface.co/athrael-soju/colqwen3.5-4.5B-v3/discussions/3#6a8876ee3a72ab808d1c3a09, here are the Sentence Transformers integrations for the Vultron retrievers, starting with Prime. I ran the NanoViDoRe v3 benchmark first, and all three line up with their ViDoRe v3 entries: Prime 0.6423 against 0.6431, Core 0.6410 against 0.6339, and Flash 0.5693 against 0.5619. The Core PR and Flash PR are both also open alongside this one, with the same changes.

Heads up, this PR was AI-generated and human-reviewed. Here's a summary of the changes as reported by my agent:

Pull Request overview

  • Integrate vultr/VultronRetrieverPrime-Qwen3.5-8B with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.
  • Add a Sentence Transformers usage section to the model card, plus the sentence-transformers and multi-vector tags.

Details

The integration is config-only: no new modeling code, and the weights are untouched. The module pipeline Transformer -> Dense -> Normalize -> MultiVectorMask mirrors colpali-engine's ColQwen3_5 forward step for step (last hidden state -> custom_text_proj -> L2 normalize -> mask by attention_mask). The custom_text_proj layer is extracted unchanged from the checkpoint into 1_Dense/. A named chat template in additional_chat_templates/sentence_transformers.jinja reproduces the ColQwen3_5_Processor prompt formats: plain text renders the query format with the ten <|endoftext|> augmentation tokens, and image inputs render the visual document format, so no manual prefixing is needed. processor_config.json now names the stock Qwen3VLProcessor so that AutoProcessor resolves without colpali-engine installed (colpali-engine constructs its own processor class explicitly, so existing usage is unaffected), and tokenizer_config.json gains padding_side: left, matching what the processor sets in code. Since the repository has no auto_map, this also makes the checkpoint usable without any extra package on top of transformers.

Verified against colpali-engine's ColQwen3_5 on text queries and image documents: the per-token embeddings match with cosine similarity 1.0000 and the MaxSim scores differ by at most 0.005.

pip install "sentence-transformers[image]>=6.0.0"
from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("vultr/VultronRetrieverPrime-Qwen3.5-8B")

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
images = [
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc1.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc2.jpg",
]

query_embeddings = model.encode_query(queries)
image_embeddings = model.encode_document(images)
print(query_embeddings[0].shape, image_embeddings[0].shape)
# torch.Size([23, 320]) torch.Size([755, 320])

# Diagonal should have higher scores
scores = model.similarity(query_embeddings, image_embeddings)
print(scores)
# tensor([[15.5859, 10.2539],
#         [ 4.9639, 12.7383]], device='cuda:0')

To try this before merging, pass revision="refs/pr/2" to MultiVectorEncoder.

Happy to tweak anything you'd like changed. Please let me know if you have any questions or feedback!

  • Tom Aarsen
tomaarsen changed pull request status to open

lgtm

Only thing that raised an eyebrow is this ViDoRe v3 entries: Prime 0.6423 against 0.6431, Core 0.6410 against 0.6339, and Flash 0.5693 against 0.5619

These seem to differ from https://mteb-leaderboard.hf.space/benchmark/ViDoRe(v3) and the difference for core is quite substantial. Does the NanoViDoRe v3 benchmark also run the private evals?

Good call, I looked into it some more and it seems that my script uses split[0]["main_score"] to grab the ViDoRe v3 scores per task directly from the MTEB results repository, instead of averaging across the subtasks per task. And no, NanoViDoRe v3 doesn't run private evals, public only. So now, with the corrected official ViDoRe v3 mean Public scores, we get this comparison:

  • Prime:
    • NanoViDoRe v3: 0.6423 NDCG@10
    • MTEB ViDoRe v3 Mean (Public): 0.6472 NDCG@10
  • Core:
    • NanoViDoRe v3: 0.6410 NDCG@10
    • MTEB ViDoRe v3 Mean (Public): 0.6372 NDCG@10
  • Flash:
    • NanoViDoRe v3: 0.5693 NDCG@10
    • MTEB ViDoRe v3 Mean (Public): 0.5649 NDCG@10

Which is rather similar on all fronts, and wouldn't really indicate to me that there's a lingering integration issue.

  • Tom Aarsen

Good call, I looked into it some more and it seems that my script uses split[0]["main_score"] to grab the ViDoRe v3 scores per task directly from the MTEB results repository, instead of averaging across the subtasks per task. And no, NanoViDoRe v3 doesn't run private evals, public only. So now, with the corrected official ViDoRe v3 mean Public scores, we get this comparison:

  • Prime:
    • NanoViDoRe v3: 0.6423 NDCG@10
    • MTEB ViDoRe v3 Mean (Public): 0.6472 NDCG@10
  • Core:
    • NanoViDoRe v3: 0.6410 NDCG@10
    • MTEB ViDoRe v3 Mean (Public): 0.6372 NDCG@10
  • Flash:
    • NanoViDoRe v3: 0.5693 NDCG@10
    • MTEB ViDoRe v3 Mean (Public): 0.5649 NDCG@10

Which is rather similar on all fronts, and wouldn't really indicate to me that there's a lingering integration issue.

  • Tom Aarsen

Thanks for checking that. split[0]["main_score"] is a tricky one. Core still seems substantially different with NanoVidore v3 vs the full Vidore V3 suite, possibly due to it being subsampled. Regardless, happy to merge this, as its integration work.

On a similar note, while working on the VultronCoder models, I put together an eval that can run the Vidore V3 eval significantly faster than the full Vidore suite, around 5 minutes vs. 1-2 hours. I left the full suite for the final eval. If you want, we can explore integrating?

possibly due to it being subsampled

Definitely, yeah. NanoViDoRe_v3 by LightOn is much smaller, it's really just for convenience to make sure that I spot any big integration issues. It's not really expected to match 1-1, and I'm actually surprised that the downscaled eval is so close to the real one.
I also check similarity between the official colpali-engine/transformers integration and the ST integration by the way, but sometimes e.g. colpali-engine updated and something broke, resulting in worse rankings. This happened a few times, so I started running NanoViDoRe to double-check.

If you want, we can explore integrating?

That does sound very cool, but I don't really tend to integrate a lot of datasets into ST, if that's what you're proposing. NanoViDoRe also isn't integrated, I just load it and initialize a MultiVectorInformationRetrievalEvaluator (https://sbert.net/docs/package_reference/multi_vector_encoder/evaluation.html#multivectorinformationretrievalevaluator). I'm wary of adding too much eval stuff to Sentence Transformers as it doesn't currently natively support e.g. fast-plaid or something, so I prefer implementing a proper evaluation in MTEB or something.

possibly due to it being subsampled

Definitely, yeah. NanoViDoRe_v3 by LightOn is much smaller, it's really just for convenience to make sure that I spot any big integration issues. It's not really expected to match 1-1, and I'm actually surprised that the downscaled eval is so close to the real one.
I also check similarity between the official colpali-engine/transformers integration and the ST integration by the way, but sometimes e.g. colpali-engine updated and something broke, resulting in worse rankings. This happened a few times, so I started running NanoViDoRe to double-check.

If you want, we can explore integrating?

That does sound very cool, but I don't really tend to integrate a lot of datasets into ST, if that's what you're proposing. NanoViDoRe also isn't integrated, I just load it and initialize a MultiVectorInformationRetrievalEvaluator (https://sbert.net/docs/package_reference/multi_vector_encoder/evaluation.html#multivectorinformationretrievalevaluator). I'm wary of adding too much eval stuff to Sentence Transformers as it doesn't currently natively support e.g. fast-plaid or something, so I prefer implementing a proper evaluation in MTEB or something.

A shame colpali-engine is getting deprecated, but less work to sync I guess.

Regarding the eval, no dataset integration needed and it's really tiny. Here's how it compares to the full mteb suite:

Vidore-fast-eval Official mteb
Scoring All on GPU: encode, upcast, MaxSim, top-k, configurable nDCG@n mode Encoding on GPU via wrapper; ranking and metrics run through mteb's CPU-side pipeline
Memory Top-k extracted without materializing the query-by-page similarity tensor Full per-chunk score matrix materialized before reduction
Precision Explicit fp32 upcast before the max-and-sum reduction; 16-bit flips near-ties Upstream scorer reduces in input dtype and casts after
Metrics Own nDCG/top-k code; writes JSON in mteb's format mteb's own implementation and results tree
eval-speed 22 tasks (Full V1 + V2 + V3 suite) in ~10 min on one B200 GPU, ~3 min on eight Hours for comparable suites, due to CPU-side ranking

But I'm conscious this PR is not the place for it. We can take it to another thread if you want.

PR Merged!

athrael-soju changed pull request status to merged
athrael-soju deleted the refs/pr/2 ref

Sign up or log in to comment