sae-splade

SAE-SPLADE (DistilBERT, k=8) — unofficial reproduction

Unofficial. This is an independent reproduction of From Tokens to Concepts: Leveraging SAE for SPLADE (Zong et al., SIGIR 2026). It was not released or endorsed by the paper's authors. Any difference from the published numbers is a property of this reproduction, not of their work.

A sparse retrieval model that replaces SPLADE's vocabulary output space with a dictionary of 65,536 learned concepts from a TopK sparse autoencoder over DistilBERT. Each token activates its 8 strongest concepts; a document is the element-wise max over its tokens. Documents end up with roughly 110 non-zero dimensions, so they stay indexable with a standard inverted index.

Results

Evaluated with the paper's protocol: full MS MARCO v1 passage collection (8.84M passages), top-1000 retrieval.

The two large-sample benchmarks match the paper closely. TREC-DL is 1.4 points low, which is within query-sampling noise for sets of 43 and 54 judged queries.

<details> <summary>Per-dataset breakdown</summary> </details>

Efficiency. The paper's central claim is comparable accuracy at roughly half of SPLADE's query-document FLOPs (0.67 vs 1.47). Our validation-time measurement converged to 0.63, consistent with that claim. A test-set FLOPs measurement is not included here yet; treat the 0.63 as indicative rather than a published number.

Usage

Download modeling_sae_splade.py from this repository (no experimaestro/xpmir needed, just torch + transformers):

from modeling_sae_splade import SAESPLADE

model = SAESPLADE.from_pretrained("Hyukkyu/sae-splade").eval()

queries = ["what causes rain"]
docs = [
    "Rain forms when water vapour in clouds condenses into droplets heavy enough to fall.",
    "The Eiffel Tower in Paris stands 330 metres tall including its antennas.",
]
print(model.score(queries, docs))     # tensor([[13.88, 0.00]])

encode_queries / encode_documents return dense tensors of width 65,536 that are ~99.8% zeros; feed the non-zero entries to an inverted index (Anserini, PISA, seismic) exactly as you would SPLADE weights.

vec = model.encode_documents([doc])[0]
nz = vec.nonzero().squeeze(-1)
postings = {int(i): float(vec[i]) for i in nz}      # concept id -> weight

Interpreting a concept

W_dec[j] is concept j's direction in DistilBERT's embedding space, which is what makes the latent space inspectable:

direction = model.W_dec[concept_id]      # (768,)

Model details

Query and document encoders share all weights; only the tokenization differs (ColBERT-style [unused0] / [unused1] markers, queries padded to 32).

Training

Two stages, both on MS MARCO v1 passages.

Final autoencoder health: reconstruction loss 17.9, 5.5% dead latents (down from a 51% peak, revived by the auxiliary loss), 15,340 distinct concepts used per batch.

Differences from the paper

Documented in full in the reproduction repository.

Limitations

• English only, trained on MS MARCO passages; out-of-domain behaviour beyond LoTTE is untested.
• Inherits DistilBERT's biases and the topical skew of MS MARCO web queries.
• Concept dimensions are not words. Unlike SPLADE, an index entry has no direct lexical reading; use W_dec to inspect one.
• fp32 only. No fp16/int8 variant is published.
• 5.5% of the dictionary is dead and contributes nothing.

License and attribution

Released under CC-BY-NC-SA-4.0, reflecting MS MARCO's non-commercial research terms. The backbone (DistilBERT) is Apache-2.0. Training used ColBERTv2 distillation scores from Stanford FutureData.

Please cite the original paper:

@article{zong2026tokens,
  title={From Tokens to Concepts: Leveraging SAE for SPLADE},
  author={Zong, Yuxuan and Vast, Mathias and Van Cooten, Basile and Soulier, Laure and Piwowarski, Benjamin},
  journal={arXiv preprint arXiv:2604.21511},
  year={2026}
}

Original authors' code: https://github.com/yzong12138/sae_splade