Let a module declare an extra output type (e.g. multi-vector + sparse from one model)
Summary
MultiVectorEncoder can already carry a second output type through the back door: a custom pipeline-final module writes an extra feature key, and encode(..., output_value=None) returns it per input (the branch at multi_vector_encoder/model.py:868-891 splits any batch-first tensor and has no key whitelist; test_encode_output_value_none_returns_feature_dicts documents that "extra keys from custom modules become user-reachable this way"). I built and validated a working model on top of this. It works, but the extra output is invisible to everything except a caller who knows to pass output_value=None. Verified against e7594f0 (6.1.0.dev0), torch 2.13.
Motivation
topk-io/Iso-ModernColBERT is a ColBERT-style model that in production also emits a sparse "SMVE" vector from the same 128-dim token embeddings (random anchor projection → per-token top-k → pooled: sum for queries, mean of non-zero contributions for documents). The sparse vector does candidate retrieval in a sparse index; exact MaxSim reranks. Users want both outputs from one forward pass, and we want to ship it as one Hub model. This is not TopK-specific: any late-interaction model with a learned or fixed sparse head, or any model that wants dense + sparse from one pass, has the same need.
What works today (the improvised route)
A custom module RandomProjectionTopK(Module) with forward_kwargs = {"task"}, appended after MultiVectorMask:
- reads
token_embeddingsandattention_mask— which, afterMultiVectorMask, is already the scoring mask (real tokens minus doc skiplist, plus expansion positions), so the sparse vector pools over exactly the tokens MaxSim scores; - writes
features["smve_embedding"]of shape(batch, width), passes everything else through; encode_query(texts, output_value=None)→ per-input dicts withtoken_embeddings,attention_mask,smve_embedding; one backbone forward; the default path andsimilarity()are unchanged (token embeddings bit-identical with/without the module);- anchors are a seeded fp32 buffer,
save/loadoverridden likeSparseAutoEncoder; a fresh process reloads it withtrust_remote_code=Trueand produces bit-identical vectors.
Sanity on SciFact: exact-MaxSim nDCG@10 0.753 (model card 0.7526); at the card's SMVE config the sparse first stage reproduces the card's recall (0.896 vs 0.900). So the mechanism is sound.
The gap
Right now the extra key is a second-class citizen:
- Undiscoverable.
encode_query/encode_documentdefault tooutput_value="token_embeddings"and hide it; users must know to passoutput_value=Noneand then slice / convert to CSR themselves. - Unscoreable.
model.similarityis MaxSim overtoken_embeddings; the sparse key has no similarity function, so evaluators,semantic_search, and thesparse_encoder/search_engines.pyintegrations (semantic_search_qdrant / _elasticsearch / _opensearch / _seismic) cannot consume it. - No hybrid hook. Candidate retrieval on one output and reranking on another has no place to plug in.
Proposal
In increasing ambition:
- Module-declared extra outputs. A module attribute such as
extra_output_keys = {"smve_embedding"}(or a registry on the model), withencode(..., output_value="smve_embedding")returning that key directly — per-input split as today, honouringconvert_to_numpy. This alone removes the "you have to know aboutNone" problem. - Per-key similarity. Let a module (or
config_sentence_transformers.json) declare asimilarity_fn_namefor its output key ("dot"for the sparse vector), somodel.similarity(..., output_value="smve_embedding")and evaluators can score it. - Hybrid retrieval helper / evaluator hook. Candidates from one output, rerank with another. The sparse search-engine helpers already index a
(batch, width)sparse matrix; they just need to be reachable from aMultiVectorEncoder.
What I'd upload to the Hub today
modules.json gaining a final {"idx": 4, "name": "4", "path": "4_RandomProjectionTopK", "type": "random_projection_topk.RandomProjectionTopK"} entry, random_projection_topk.py at the repo root, and 4_RandomProjectionTopK/{config.json, model.safetensors}. Loads with MultiVectorEncoder(repo, trust_remote_code=True).
Happy to send a sketch PR for (1)+(2), and to share the module and eval scripts (the module is generic — random projection, top-k, pooling — nothing vendor-specific).
Source: huggingface/sentence-transformers