Sparse Retrieval, Part 2: SPLADE
How SPLADE expands and reweights terms while keeping an inverted-index retrieval path.
Moving on from BM25 (Part 1). BM25 searches for almost the same words: the exact words, or, if the preprocessing does stemming (covered in Part 1), "running" and "runs" for the query "run". But it can't find a doc about a "marathon" when the query is "run". Because how will it? "marathon" is just a different word, it has its own posting list, and the query never touches it.
So either we define a dictionary mapping each word to its synonyms, and search for them too. That would be the brute-force way, and it must have come to the minds of Searchacians (it did). But who writes this dictionary, for every word, for every domain? And "run" next to "python" is not "marathon" anymore.
But no, the field of ML, DL solves it. How? SPLADE. What is SPLADE? -> What I had read about it at first: it gives you a vector for the given doc, a very high-dimensional vector, which stores mostly 0s, and the index at which the value is not 0 is a token ID, with a weight there, like 1.7. So it also has a vocabulary mapping every token to a token ID. Now when we search, related words match too, even if they were never in the doc. Still vague, so in this post we open it up: from that black box, down to BERT's insides (tokens, embeddings, attention, the MLM head), and back up to how SPLADE builds these vectors, how it's trained, and how search works with them.
Contents (the 13 min covers the main text; the dotted links open optional side notes, which add more on top):
- Two roads
- SPLADE, as a black box
- Down the rabbit hole (optional): BERT from scratch
- Climbing back out: the SPLADE formula
- BERT knows words, not relevance
- Searching with it
- What SPLADE doesn't see
- What's next
Two roads
Before going in: there are two roads here, and I've already been down the first one.
- Dense. In the typical single-vector setup, turn every doc into one vector of ~768 numbers, all of them non-zero, and for a query, find the docs with the nearest vectors. It does get "run" ≈ "marathon". But instead of the inverted index, you need a different kind of index for nearest-neighbour search (that's the HNSW index I built for VortexDB: the mental model, the Rust implementation), exact matches like a product ID or a person's name get fuzzy, and you can't really tell why a doc matched. (more on dense)
- Learned sparse. Keep the inverted index from Part 1, keep the speed, but let a model decide which words a doc should be indexed under, and with what weight.
SPLADE is the second road, and the one this series is about.
SPLADE, as a black box
So, SParse Lexical AnD Expansion model.
For now, we don't look inside. All we need to know: SPLADE is built on BERT, a model trained on huge amounts of text, which has learned the associations of run and marathon. So BERT knows that run and marathon are related. So SPLADE's first idea is simple: use BERT to expand a document with related words, like marathon and jogging for run. The model quietly adds weights to these words too, even when they never appeared in the text. This is the Expansion part.
Next thing that SPLADE does: BERT's output is very dense, a huge list of numbers, hundreds of dimensions, nothing like our sparse, word-based index. But SPLADE forces the output into a different shape: instead of hundreds of dense numbers, the output has one slot for every token in BERT's vocabulary, 30,522 slots. For each slot, the model predicts a weight: how relevant that word is to this document.
At first, thousands of these values are small but non-zero. SPLADE adds a penalty during training that pushes these weights down, and anything the model scores at or below 0 becomes exactly 0. So only a small number of words keep a meaningful weight: with the settings in the papers, typically a few dozen to a few hundred for a passage, fewer for a short query, but the count depends on how hard the penalty is set. This penalty is called sparse regularization.
So this is what we're going for:
Now I will open this black box a little, to explain the above points better.
So BERT returns an embedding for every token of the document, 768 numbers per token, none of them 0. Now these numbers, they don't correspond to anything we understand. It's not like the 368th number is meant for the word "run", it's just an abstract dimension in maths space, meaningless to us, and more importantly, to an inverted index, which needs actual tokens/words, not these abstract dimensions.
Now SPLADE reshapes this thing. It forces the output to have exactly one slot per vocabulary token, all 30,522 tokens BERT knows. For example, slot 1037 is "a", slot 2448 is "run", and slot 8589 is "marathon" (these are the real IDs). So in this, every word gets its own fixed, known position.
So if SPLADE only does this reshaping, you will get 30,522 numbers, and thousands of them would be small but non-zero. Why so many? Because the layer SPLADE reads from gives a score to every vocab token by design, and plenty of those come out positive, we'll see it in a bit. That's not useful for an inverted index, whose whole purpose is that you only store and look up words that actually matter for a document. If every document had a small weight on all 30,522 words, your posting lists would include every document for every word, defeating the entire purpose of fast lookup.
So that's where the sparsity regularization comes in, the penalty during training. It pushes the weights down, hardest on words that show up in lots of docs, so the model learns to score most words at or below 0. The exact zeros come from a ReLU we'll see later, which turns every such score into 0. The penalty doesn't zero anything by itself or fix how many words survive, it just makes keeping a word cost something. So for our "run" example, after training, you'd end up with maybe "run" at a high weight, "marathon" and "jogging" at smaller but real weights, and almost all the other vocabulary words sitting at exactly zero, not stored.
How does it reshape? It uses BERT's internals, one of its layers. Here's where we're headed, the whole of SPLADE in one line:
is BERT's score for vocab word at token of the doc, and is the doc's final weight for word . Every piece of this gets unpacked below.
Down the rabbit hole (optional)
Already know transformers? Skip to the SPLADE formula. Otherwise, read on.
Okay, this is where I went down a rabbit hole. The SPLADE paper says it uses the "logits of the MLM layer". What's MLM? It's how BERT was trained. What's BERT? The encoder half of a transformer. What's a transformer? Attention, stacked. And attention needs Q, K, V, from where?? And the word vectors themselves, from where?? It kept going:
So now we start at the bottom of the hole and build back up, step by step, from raw text all the way up to the layer SPLADE uses, and I'll try to explain almost everything, assuming you know nothing about transformers (I didn't).
Step 1: text → tokens → IDs
A model only does maths, so first, text has to become numbers. BERT has a fixed vocabulary of 30,522 tokens. The tokenizer (it's called WordPiece) splits the text into tokens from that list and replaces each one with its ID, its row number in the list. Words that are not in the list get broken into pieces it does know: "splade" becomes sp, ##lad, ##e (the ## means "glued to the previous piece"). Funny enough, "jog" is not a token, it's jo + ##g, but "jogging" is one.
Remember this vocabulary, because SPLADE's 30,522 slots are exactly these tokens. SPLADE's "words" are BERT's tokens.
Step 2: IDs → vectors
But an ID is just a label. 2448 isn't "bigger" than 1037 in any useful way. So BERT has a big table called the embedding matrix , one row per vocab token: 30,522 rows × 768 numbers. The vector for a token is just its row. That's it, a lookup:
| d1 | d2 | d3 | d4 | |
|---|---|---|---|---|
| 0 the | 0.10 | 0.20 | 0.10 | 0.10 |
| 1 river | 2.00 | 0.40 | 0.00 | 1.00 |
| 2 bank | 1.20 | 1.60 | 0.60 | 0.20 |
| 3 flows | 1.80 | 0.00 | 1.40 | 0.40 |
| 4 money | 0.30 | 1.90 | 0.20 | 0.10 |
| 5 loan | 0.20 | 1.70 | 0.50 | 0.00 |
| 1.20 | 1.60 | 0.60 | 0.20 |
Where did these numbers come from? Nobody wrote them, they start random and get learned during training (how). One more thing: before going in, every token's vector also gets a position vector added to it, so the model knows the word order (and a segment vector in BERT, saying which sentence it belongs to).
Step 3: one vector per word, but "bank"?
Now, same token, same row, always. "bank" in "river bank" and "bank loan" gets the exact same vector. That's not really better than BM25. Each word needs to look at its neighbours and change itself based on them. That's attention.
Every token makes three new vectors out of its own vector, by multiplying it with three matrices:
- a query : what am I looking for?
- a key : what do I contain? like a label on me
- a value : what do I hand over if someone picks me?
Then for one word, we compare its query with every word's key using the dot product (big when the two vectors point the same way, and also when they're long: unlike cosine similarity, it doesn't normalize the lengths), divide by to keep the numbers from getting too big, softmax them into percentages that add up to 100%, and the word's new vector is the weighted mix of everyone's values. For all words at once:
, , are just all the words' queries, keys and values stacked as rows, is the length of each key (64 in BERT), and is every query dotted with every key.
| river | bank | flows | |
|---|---|---|---|
| river | 0.36 | 0.50 | 0.14 |
| bank | 0.77 | 0.13 | 0.10 |
| flows | 0.48 | 0.27 | 0.25 |
See it step by step, with every number. And yes, , , are learned too, same as the embeddings. BERT does this 12 times in parallel in every layer, each with its own matrices (12 "heads" × 64 numbers = 768).
Step 4: the rest of a layer, and stacking them
After attention, every token goes through a small neural network on its own (768 → 3,072 → 768, with a GELU in between), plus some plumbing that keeps things stable. That's one layer. BERT-base stacks 12 of them.
What comes out of the top: one 768-number vector per token, that now "knows" its context. Not one per document, one per token. The "bank" next to "river" now has a different vector from the "bank" next to "loan".
(BERT is only one half of the original transformer, by the way. Transformer vs BERT vs GPT.)
Step 5: how BERT learned: hide a word, guess it
BERT was trained on Wikipedia and a pile of books, ~3.3 billion words, with a game: hide 15% of the tokens (mostly behind a [MASK]) and guess them, using the words on both sides. This is the Masked Language Model, MLM. The guessing is done by a small layer on top, the MLM head: it takes a token's final 768-number vector and gives a score (a "logit") for every one of the 30,522 vocab tokens. Highest score = BERT's guess.
And this is it. This is the layer SPLADE uses.
Think about it: for the token "run" in "i went for a run", the MLM head scores all 30,522 vocab tokens by "how well would this token fit here". Something like (illustrative, not a real model output): run high, jogging high, marathon some, banana very low. That's literally an expansion, already there, BERT learned it for free while playing its guessing game. And that's the "score for every vocab token" from earlier, the reason we get thousands of non-zeros before the penalty. (what's inside the MLM head)
Climbing back out: the SPLADE formula
Okay, back up. Now we have, for every token in the doc, a score for every vocab token . A grid: (tokens in the doc) × 30,522. SPLADE turns this grid into one vector for the doc:
Three moves, read it from inside out:
- ReLU: negatives become 0. BERT says "banana" doesn't fit here at all, fine, that's a 0, not a negative.
- log(1 + x): saturation. Remember BM25's saturation? Same idea: the 10th unit of confidence should add less than the 1st, so one word can't dominate just because BERT is very sure about it. And , so zeros stay zeros (no log(0) problem here).
- max over tokens: for every vocab word, keep its best score from any token in the doc.
| run | ran | marathon | race | running | she | a | banana | |
|---|---|---|---|---|---|---|---|---|
| she | -1.20 | 0.30 | -0.80 | -1.50 | -1.00 | 4.10 | -0.50 | -3.00 |
| ran | 3.20 | 5.00 | 0.90 | 1.40 | 2.10 | -0.40 | -1.10 | -2.60 |
| a | 0.40 | -0.20 | 0.60 | 0.80 | -0.60 | -1.00 | 1.20 | -2.20 |
| marathon | 2.00 | 0.50 | 6.00 | 2.60 | 1.70 | -1.30 | -0.90 | -2.90 |
(The first SPLADE paper actually used a sum over tokens here, the second one switched to max. Why.)
BERT knows words, not relevance
One thing I missed at first: BERT alone knows run and marathon are related, but it doesn't know what makes a doc relevant to a query. Those are different. So SPLADE starts from BERT's weights and is trained more (fine-tuned) on MS MARCO, a dataset of real Bing search queries with passages marked as relevant.
For each query, we take its relevant passage , a passage BM25 likes but that isn't relevant (a "hard negative"), and the other queries' passages in the same batch as extra negatives. Then we push the relevant one to the top:
If that looks familiar, it's the same softmax from attention, just over docs this time: "what % goes to the right doc", and we want it near 100%.
And then the penalty, added on top, with two knobs, one for queries and one for docs:
The penalty is called the FLOPS regularizer, and the fun thing about it: it doesn't just push all weights down, it pushes hardest on words that are active in lots of docs. Which are exactly the words with long posting lists, the ones that make search slow. Bigger , fewer words survive:
Why two knobs? The query is what we pay for at search time, every time, so you can squeeze queries harder than docs.
Searching with it
Now putting it all together.
Indexing: run every doc through SPLADE once, offline (a GPU makes this much faster, but it isn't required), and for each non-zero (token, weight), put the doc in that token's posting list with the weight. It's the same inverted index from Part 1, just storing learned weights instead of TF counts.
Search: run the query through SPLADE too, so the query gets expanded too, and the score is just a dot product over the tokens both have:
And because both vectors are mostly 0, we only need to walk the posting lists of the query's few non-zero tokens. Same trick as BM25.
| doc | BM25 matches | SPLADE score |
|---|---|---|
| A: she ran a marathon | nothing, score 0 | 4.67 |
| B: how to run a python script | run | 2.73 |
| C: jogging every morning | morning | 9.21 |
There's also a cheaper version, SPLADE-doc, that doesn't run BERT on the query at all.
And how much better is it? On MS MARCO, with MRR@10 (Mean Reciprocal Rank: for each query, find the first relevant passage in the top 10: rank 1 scores 1, rank 2 scores , rank 3 scores , not in the top 10 scores 0. Average over all queries, . Higher is better), on the MS MARCO dev set, from the SPLADE v2 paper (it reports them as fractions like 0.368, here ×100):
| model | MRR@10 |
|---|---|
| BM25 | 18.4 |
| doc2query-T5 | 27.7 |
| SPLADE (v1) | 32.2 |
| SPLADE-max | 34.0 |
| DistilSPLADE-max | 36.8 |
The best one gets about twice BM25's MRR@10, while still being an inverted index.
One detail: everything in the rabbit hole above is BERT-base, but the v2 models in this table start from DistilBERT-base, a smaller 6-layer version of BERT with the same vocabulary and the same kind of MLM head.
What SPLADE doesn't see
Not free, though:
- BERT everywhere. Every doc goes through BERT at index time, and every query at search time (unless SPLADE-doc). BM25 needs no model at all.
- Length limit. BERT takes at most 512 tokens, and SPLADE was trained on 256, so long docs have to be split into chunks.
- BERT's vocabulary. The index stores BERT's tokens. A new product name or a typo gets chopped into pieces like
##lad, and the weights land on those pieces. - Expansion can be wrong. The model can add words that don't fit the doc, and it can't tell you why it did.
- Bigger index. Expansion terms make posting lists longer than BM25's. The FLOPS penalty keeps this in check, but it's a trade-off you tune.
What's next
So that's SPLADE: BERT's guessing game, reshaped into an inverted index. I've been adding sparse vector support to VortexDB, so next I'll actually build it: storing these sparse vectors, an inverted index with weights, and searching it fast. Stay tuned.
References
- SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking (Formal, Piwowarski, Clinchant, SIGIR 2021)
- SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval (Formal, Lassance, Piwowarski, Clinchant, 2021)
- BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin, Chang, Lee, Toutanova, NAACL 2019)
- Minimizing FLOPs to Learn Efficient Sparse Representations (Paria, Yeh, Yen, Xu, Ravikumar, Póczos, ICLR 2020)
- The Illustrated Transformer (Jay Alammar)