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.

Benchmark	Metric	This model	Paper	Δ
MS MARCO dev small (6,980 q)	MRR@10	37.68	37.6	+0.08
LoTTE search, 5 domains	Success@5	68.66	68.8	−0.14
TREC-DL 2019 + 2020	nDCG@10	70.68	72.1	−1.42

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.

Per-dataset breakdown

Dataset	Metric	Score
MS MARCO dev small	MRR@10	37.68
TREC-DL 2019	nDCG@10	70.40
TREC-DL 2020	nDCG@10	70.96
LoTTE Writing	Success@5	76.00
LoTTE Recreation	Success@5	68.40
LoTTE Science	Success@5	54.13
LoTTE Technology	Success@5	63.09
LoTTE Lifestyle	Success@5	81.69

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


Backbone	`distilbert-base-uncased`, 6 layers, last hidden state
Concept dictionary	65,536 latents
Active concepts per token (k)	8
Parameters	167.2M total — 66.4M backbone, 100.8M autoencoder
Precision	fp32
Query / document truncation	32 / 256 tokens
Pooling	max over tokens, then `log1p`, scaled by a learned `alpha` (1.18)
Scoring	dot product

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.

	Stage 1 — SAE pretraining	Stage 2 — SAE-SPLADE fine-tuning
Objective	reconstruction MSE + auxiliary dead-latent loss (0.0625)	KL (1.0) + MarginMSE (0.05) distillation + FLOPS regularization
Teacher	—	ColBERTv2 64-way scores, 8-way sampled
Backbone	frozen	trained
Steps	160,000	240,000
Batch	768 documents	32 queries × 8 documents
Optimizer	AdamW, wd 0.01, eps 6e-10	AdamW, wd 0.01, eps 1e-8
Learning rate	5e-5	2e-5
Schedule	10k warmup, then near-constant	10k warmup, then linear decay to 0
FLOPS λ (doc / query)	—	0.04 / 0.06, quadratic warmup over 6k steps

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.

Aspect	Paper	Here	Effect
Hardware	1× A100	2× TITAN RTX, data-parallel	none by design: gradients are averaged so the update equals the single-GPU global batch (verified, max relative gradient error 1.5e-7)
Micro-batching	OOM-driven adaptive split	fixed split (64 docs / 8 queries per GPU)	needed for deterministic collectives; changes nothing mathematically
FLOPS grouping	per adaptive micro-batch (size unreported)	per 16 gathered queries	the one setting that could shift sparsity, since FLOPS squares batch means
Compute	~35h + ~24h	38h + 47h + ~7h eval	—

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