Let a module declare an extra output type (e.g. multi-vector + sparse from one model)

Author: davidmyrielCreated Aug 18, 2026Updated Sep 4, 2026

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_embeddings and attention_mask — which, after MultiVectorMask, 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 with token_embeddings, attention_mask, smve_embedding; one backbone forward; the default path and similarity() are unchanged (token embeddings bit-identical with/without the module);
  • anchors are a seeded fp32 buffer, save/load overridden like SparseAutoEncoder; a fresh process reloads it with trust_remote_code=True and 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_document default to output_value="token_embeddings" and hide it; users must know to pass output_value=None and then slice / convert to CSR themselves.
  • Unscoreable. model.similarity is MaxSim over token_embeddings; the sparse key has no similarity function, so evaluators, semantic_search, and the sparse_encoder/search_engines.py integrations (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:

  1. Module-declared extra outputs. A module attribute such as extra_output_keys = {"smve_embedding"} (or a registry on the model), with encode(..., output_value="smve_embedding") returning that key directly — per-input split as today, honouring convert_to_numpy. This alone removes the "you have to know about None" problem.
  2. Per-key similarity. Let a module (or config_sentence_transformers.json) declare a similarity_fn_name for its output key ("dot" for the sparse vector), so model.similarity(..., output_value="smve_embedding") and evaluators can score it.
  3. 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 a MultiVectorEncoder.

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