← Writing
SearchInformation RetrievalSPLADE

Sparse Retrieval, Part 2: SPLADE

How SPLADE expands and reweights terms while keeping an inverted-index retrieval path.

Oct 202613 min read

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

Before going in: there are two roads here, and I've already been down the first one.

  1. 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)
  2. 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:

doc: “i went for a run”
BM25, term frequency: 5 of 30,522 slots non-zero
[ 0…0, 1,a0…0, 1,i0…0, 1,for0…0, 1,went0…0, 1,run0…0 ] 
SPLADE, learned weights: 7 of 30,522 slots non-zero
[ 0…0, 0.6,went0…0, 2.4,run0…0, 0.5,+ran0…0, 1.5,+running0…0, 0.8,+exercise0…0, 0.9,+marathon0…0, 1.2,+jogging0…0 ] 
Same doc, two sparse vectors. BM25 counts the words that are there, SPLADE weighs them and adds new ones.

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:

wj=max⁡i∈doclog⁡(1+ReLU(wij))w_j = \max_{i \in \text{doc}} \log\big(1 + \text{ReLU}(w_{ij})\big)

wijw_{ij} is BERT's score for vocab word jj at token ii of the doc, and wjw_j is the doc's final weight for word jj. 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:

SPLADE: “uses the logits of BERT's MLM layer”
↳ what's MLM? → how BERT was trained
↳ what's BERT? → the encoder half of a transformer
↳ what's a transformer? → attention, stacked
↳ attention needs Q, K, V. from where??
↳ and the word vectors themselves, from where??
↳ …and why is there a GELU here
↑ climbing back out to SPLADE
How it went. Each answer had another question in it.

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.

[CLS]
101
i
1045
went
2253
for
2005
a
1037
run
2448
[SEP]
102
Real BERT (bert-base-uncased) tokens and IDs. [CLS] and [SEP] are special tokens BERT adds at the start and end.

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 EE, one row per vocab token: 30,522 rows × 768 numbers. The vector for a token is just its row. That's it, a lookup:

xbank=Eid(bank)E∈R30522×768\mathbf{x}_{\text{bank}} = E_{\text{id(bank)}} \qquad E \in \mathbb{R}^{30522 \times 768}
d1d2d3d4
0 the0.100.200.100.10
1 river2.000.400.001.00
2 bank1.201.600.600.20
3 flows1.800.001.400.40
4 money0.301.900.200.10
5 loan0.201.700.500.00
→
xbank=E2,:\mathbf x_{\text{bank}} = E_{2,:}
1.201.600.600.20
A toy version: 6 words, 4 numbers each.

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 q=xWQ\mathbf{q} = \mathbf{x} W^Q: what am I looking for?
  • a key k=xWK\mathbf{k} = \mathbf{x} W^K: what do I contain? like a label on me
  • a value v=xWV\mathbf{v} = \mathbf{x} W^V: 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 dk\sqrt{d_k} 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:

Attention(Q,K,V)=softmax(QK⊤dk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) V

QQ, KK, VV are just all the words' queries, keys and values stacked as rows, dkd_k is the length of each key (64 in BERT), and QK⊤QK^\top is every query dotted with every key.

riverbankflows
river0.360.500.14
bank0.770.130.10
flows0.480.270.25
zbank=0.77 vriver+0.13 vbank+0.10 vflows\mathbf z_{\text{bank}} = 0.77\,\mathbf v_{\text{river}} + 0.13\,\mathbf v_{\text{bank}} + 0.10\,\mathbf v_{\text{flows}}
=[1.24, 0.69]= [1.24,\ 0.69]
A toy "river bank flows" with 4-number word vectors. Rows = the word asking, columns = who it listens to. "bank" puts 77% of its attention on "river", so its new vector is mostly river-flavoured. Numbers are a toy, but computed for real.

See it step by step, with every number. And yes, WQW^Q, WKW^K, WVW^V 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.

to the next layerONE ENCODER LAYER×12Multi-head attentionAdd & NormFFNFFNAdd & Normx₁riverx₂bank+ position
One layer, two tokens. Attention mixes the tokens, the FFN works on each token alone, and the dashed lines carry each block's input around it to be added back. The output goes into the next layer, 12 times. Drawn after the Illustrated Transformer.

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.

run
0.31
walk
0.24
swim
0.08
drive
0.07
ride
0.05
banana
0.00
top guesses out of all 30,522, after softmax. made-up but realistic numbers
What the MLM head says for a masked word. Same word "bank", different neighbours, different guesses. Illustrative numbers.

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 ii in the doc, a score wijw_{ij} for every vocab token jj. A grid: (tokens in the doc) × 30,522. SPLADE turns this grid into one vector for the doc:

wj=max⁡i∈doclog⁡(1+ReLU(wij))w_j = \max_{i \in \text{doc}} \log\big(1 + \text{ReLU}(w_{ij})\big)

Three moves, read it from inside out:

  1. ReLU: negatives become 0. BERT says "banana" doesn't fit here at all, fine, that's a 0, not a negative.
  2. 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 log⁡(1+0)=0\log(1 + 0) = 0, so zeros stay zeros (no log(0) problem here).
  3. max over tokens: for every vocab word, keep its best score from any token in the doc.
02468-202468logit wᵢⱼweight
log(1 + ReLU(x))ReLU(x)
The ReLU cuts the negatives, the log squashes the big ones.
wijw_{ij}
runranmarathonracerunningsheabanana
she-1.200.30-0.80-1.50-1.004.10-0.50-3.00
ran3.205.000.901.402.10-0.40-1.10-2.60
a0.40-0.200.600.80-0.60-1.001.20-2.20
marathon2.000.506.002.601.70-1.30-0.90-2.90
rows = tokens of the doc, columns = 8 of the 30,522 vocab slots.
A 4-token doc and 8 of the 30,522 columns. Note "marathon" in the doc also lights up "race" and "run", and "ran" lights up "running". Also note "a" and "she" survive with real weight. That's what the penalty in the next section is for. Logits are made up, the maths is real.

(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 d+d^+, a passage BM25 likes but that isn't relevant d−d^- (a "hard negative"), and the other queries' passages in the same batch as extra negatives. Then we push the relevant one to the top:

Lrank=−log⁡es(q,d+)es(q,d+)+es(q,d−)+∑jes(q,dj−)\mathcal{L}_{\text{rank}} = -\log \frac{e^{s(q, d^+)}}{e^{s(q, d^+)} + e^{s(q, d^-)} + \sum_j e^{s(q, d_j^-)}}

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:

L=Lrank+λq Lregq+λd Lregd\mathcal{L} = \mathcal{L}_{\text{rank}} + \lambda_q \, \mathcal{L}^{q}_{\text{reg}} + \lambda_d \, \mathcal{L}^{d}_{\text{reg}}

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 λ\lambda, fewer words survive:

marathon
1.95
ran
1.79
she
1.63
run
1.44
race
1.28
running
1.13
a
0.79
jogging
0.50
athlete
0.40
finish
0.30
A toy of the effect, not the real training. As λ goes up, "a" dies first (it's in almost every doc), then the weak expansions, and the words that actually describe the doc stay.

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:

s(q,d)=∑jqj djs(q, d) = \sum_{j} q_j \, d_j

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.

indexing, once, offline
doc→SPLADE (BERT + MLM head + max)→{marathon 1.95, ran 1.79, run 1.44, …}→posting lists: marathon → (A, 1.95) …
search, every query
query→SPLADE (or just its tokens)→{run 2.1, morning 1.8, …}→walk those posting lists, add up qⱼ × dⱼ→ranked docs
The whole thing, end to end.
query “morning run” → run 2.1 · morning 1.8 · running 1.1 · jogging 0.8 · exercise 0.5
docBM25 matchesSPLADE score
A: she ran a marathonnothing, score 04.67
B: how to run a python scriptrun2.73
C: jogging every morningmorning9.21
doc A: marathon 1.95 · ran 1.79 · run 1.44 · race 1.28 · running 1.13 · jogging 0.5
s(q,dA)=2.1×1.44⏟run+1.1×1.13⏟running+0.8×0.5⏟jogging=4.67s(q, d_{A}) = \underbrace{2.1 \times 1.44}_{\text{run}} + \underbrace{1.1 \times 1.13}_{\text{running}} + \underbrace{0.8 \times 0.5}_{\text{jogging}} = 4.67
Query "morning run" against three docs. BM25 can never find A (no shared word) and happily matches B on "run". SPLADE ranks C, then A, then B. Weights are illustrative.

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 12\frac{1}{2}, rank 3 scores 13\frac{1}{3}, not in the top 10 scores 0. Average over all queries, ×100\times 100. Higher is better), on the MS MARCO dev set, from the SPLADE v2 paper (it reports them as fractions like 0.368, here ×100):

modelMRR@10
BM2518.4
doc2query-T527.7
SPLADE (v1)32.2
SPLADE-max34.0
DistilSPLADE-max36.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

  1. SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking (Formal, Piwowarski, Clinchant, SIGIR 2021)
  2. SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval (Formal, Lassance, Piwowarski, Clinchant, 2021)
  3. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin, Chang, Lee, Toutanova, NAACL 2019)
  4. Minimizing FLOPs to Learn Efficient Sparse Representations (Paria, Yeh, Yen, Xu, Ravikumar, Póczos, ICLR 2020)
  5. The Illustrated Transformer (Jay Alammar)