Learned sparse retrieval with SPLADE

This uses distilbert-base-uncased's existing masked-language-model head to demonstrate the SPLADE mechanism directly — this is not a real, specially-trained SPLADE checkpoint. A production SPLADE model is fine-tuned with a contrastive ranking loss plus explicit sparsity regularisation; an off-the-shelf, un-fine-tuned MLM head produces the same shape of output without the same retrieval-quality guarantees. Treat the term weights below as illustrative, not as calibrated relevance scores.

Setup#

pip install transformers torch

distilbert-base-uncased downloads once, roughly 260 MB. Outputs verified with transformers 5.6.2 and torch 2.5.1 on CPU.

Building a sparse vector, and finding query-document overlap#

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased")
model.eval()

SKIP = set(tokenizer.all_special_tokens) | {".", ",", ";", "?", "!", "with", "and", "the", "a", "of", "for"}


def splade_vector(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits[0]              # (seq_len, vocab_size)
    weights = torch.log1p(torch.relu(logits))            # log(1 + relu(x)) per token per vocab entry
    sparse, _ = weights.max(dim=0)                        # max-pool over sequence positions
    return sparse


def top_terms(vec, k=8):
    values, idx = vec.topk(k + len(SKIP))
    out = []
    for v, i in zip(values, idx):
        tok = tokenizer.convert_ids_to_tokens([i.item()])[0]
        if tok in SKIP:
            continue
        out.append((tok, round(v.item(), 2)))
        if len(out) == k:
            break
    return out


doc = "affordable running shoes with good cushioning"
vec = splade_vector(doc)
print("input text:", repr(doc))
print("non-zero dimensions:", int((vec > 0).sum()), "out of", vec.numel(), "vocabulary entries")
print("top activated terms (includes words not in the original text):")
for term, weight in top_terms(vec):
    print(f"  {term:<12} {weight}")

query = "cheap jogging sneakers"
q_vec = splade_vector(query)
overlap = torch.minimum(vec, q_vec)
nz = overlap.nonzero().squeeze(-1)
contribs = [(tokenizer.convert_ids_to_tokens([i.item()])[0], round(overlap[i].item(), 2)) for i in nz]
contribs = [(t, c) for t, c in contribs if t not in SKIP]

print(f"\nquery: {query!r}")
print("shared activated terms and their contribution to the match score:")
for term, c in sorted(contribs, key=lambda x: -x[1])[:8]:
    print(f"  {term:<12} {c}")
print(f"sparse dot-product score: {overlap.sum().item():.2f}")

input text: 'affordable running shoes with good cushioning'
non-zero dimensions: 6470 out of 30522 vocabulary entries
top activated terms (includes words not in the original text):
  shoes        2.96
  running      2.9
  good         2.88
  cushion      2.88
  ##ing        2.75
  feet         2.48
  walking      2.46
  affordable   2.45

query: 'cheap jogging sneakers'
shared activated terms and their contribution to the match score:
  shoes        2.75
  sandals      2.43
  walking      2.42
  boots        2.4
  sneakers     2.33
  running      2.31
  shoe         2.31
  barefoot     2.17
sparse dot-product score: 2162.54

Neither "sandals," "boots" nor "barefoot" appear anywhere in the original document text, yet all three activated strongly enough to contribute to the match score against a query about "sneakers" — the model expanded the document's vocabulary using its own learned sense of which words belong to the same world as "shoes" and "running."

The walkthrough#

torch.log1p(torch.relu(logits)) is the whole SPLADE transform. relu zeroes out every negative logit (words the model actively thinks are unlikely given this context), and log1p compresses the remaining positive values so a handful of very confident words cannot completely dominate the vector.

weights.max(dim=0) pools across token positions. A word can be strongly implied by any position in the sequence — max-pooling takes each vocabulary word's single strongest activation across the whole input, rather than averaging it down across positions where it was less relevant.

torch.minimum(vec, q_vec) computes term-by-term overlap. For every vocabulary word, this takes whichever of the query's or document's activation is smaller — a word only contributes to the match score in proportion to how strongly both sides activated it, which is the sparse-vector equivalent of a dot product when most dimensions are zero on at least one side.

SKIP filters punctuation and a handful of stop words for a readable demo. A real, trained SPLADE model suppresses these itself, as part of its training objective — shown here as a manual filter purely to keep the printed output legible.

Common mistakes#

Treating this untrained-for-the-task MLM head as retrieval-quality output. As stated up front, this demonstrates the computation, not calibrated relevance. A real SPLADE checkpoint is trained with a ranking loss over labelled query-document pairs specifically so its activations reflect genuine relevance, not only generic language-model plausibility.

Ignoring how many dimensions activate. 6,470 out of 30,522 vocabulary entries activated for one six-word sentence here — far denser than a typical trained SPLADE vector, which is explicitly regularised toward much greater sparsity during training. This matters directly for storage and search cost at scale.

Comparing the sparse dot-product score across different document lengths without normalisation. Longer documents naturally activate more terms and accumulate a higher raw score, independent of relevance — production sparse-retrieval systems normalise for this.

Try it yourself#

Change the document to a completely different sentence, such as "a laptop with a fast processor and long battery life", and re-run the query "cheap jogging sneakers" against it. Confirm the shared-terms overlap collapses toward nothing — a useful sanity check that the expansion mechanism responds to genuine topical relatedness, not to any input at all.

What to learn next#

• Hybrid search with reciprocal rank fusion — combining a sparse signal like this with a dense one.
• ColBERT and late interaction — a different way to keep search word-aware instead of relying on one blurred vector.
• Building a retrieval test set for your own corpus — properly measuring whether a method like this helps on your own data.