Gen AI

Learned Sparse Retrieval

Learned sparse retrieval uses a neural network to inject context and synonyms into a lexical search index, creating a best-of-both-worlds search strategy.

Unlike dense embeddings which output a short, thick vector, learned sparse models output a vocabulary-sized vector consisting mostly of zeros, expanding documents with related terms.

Why Does This Exist?

In the world of search and retrieval, there are two primary paradigms:

Lexical Search (e.g., BM25): Fast, exact, and explainable, but it utterly fails if the query and the document use different synonyms (e.g., "automobile" vs "car").
Dense Semantic Search (e.g., BERT embeddings): Excellent at understanding meaning and synonyms, but computationally expensive, unexplainable, and prone to hallucinating connections (lacking exact keyword precision).

For a long time, the solution was simply to run both independently and mathematically merge their scores (Hybrid Search).

However, a newer approach called Learned Sparse Retrieval (with the most famous algorithm being SPLADE) attempts to merge the benefits of both paradigms into a single index. It uses a transformer neural network to understand semantic meaning, but instead of outputting a dense block of numbers, it outputs a sparse vector that maps directly to the vocabulary. It effectively performs "document expansion" during the indexing phase, allowing you to use a traditional, blazing-fast inverted index while still getting the semantic understanding of a neural network.

Think of It Like This

The overly helpful librarian

Imagine a standard inverted index (BM25) as a literal-minded librarian. You hand them a book titled "Canine Health," and they write down exactly two index cards: one for "Canine" and one for "Health." If someone searches for "dog," the librarian will confidently say they have no books on that topic.

Now imagine a Learned Sparse model as an overly helpful, highly-educated librarian. You hand them the same book, "Canine Health." The librarian reads it, understands the semantic meaning, and creates index cards not just for the literal words in the title, but for related concepts. They create cards for "Canine," "Health," "Dog," "Puppy," "Vet," and "Veterinarian." They also assign a weight to each card based on how relevant it is.

When a user comes in searching for "dog," the librarian immediately pulls up the "Canine Health" book because they proactively indexed the synonym. The search is still a fast, exact index lookup, but the index itself has been semantically expanded.

How It Actually Works

The Dense vs. Sparse Vector

A standard dense embedding model maps a document to a vector of fixed size (e.g., 768 dimensions), where every number is a non-zero float. The dimensions do not correspond to human-readable words; they are latent concepts.

A learned sparse model maps a document to a vector where the number of dimensions is equal to the size of the entire tokenizer vocabulary (e.g., 30,522 dimensions for BERT). Crucially, almost all of these 30,522 numbers are exactly zero. Only a small handful (perhaps 100 or 200) are non-zero.

Document Expansion (SPLADE)

The most prominent architecture for this is SPLADE (Sparse Lexical and Expansion Model for First Stage Retrieval). Here is how it processes a document during indexing:

Forward Pass: The document is fed through a transformer (like BERT).
MLM Head: Instead of pooling the output into a single dense vector, SPLADE uses the Masked Language Modeling (MLM) head. This is the part of BERT normally used to predict missing words in a sentence. SPLADE asks the MLM head to predict all the relevant words for the document across the entire vocabulary.
Max Pooling: It aggregates these predictions across all the tokens in the document.
Sparsification: It applies a regularization penalty (like L1 regularization) to force the network to push the vast majority of the vocabulary weights to exactly zero.

The output is a dictionary mapping specific vocabulary tokens to importance weights. For a document containing the sentence "The quick brown fox", the SPLADE output might look like: {"quick": 1.2, "brown": 0.8, "fox": 2.1, "fast": 0.9, "animal": 0.5, "wildlife": 0.3}

Notice that "fast", "animal", and "wildlife" were not in the original text. The neural network expanded the document to include them.

Inverted Index Compatibility

Because the output is just a list of words and their weights, it can be dumped directly into a traditional Inverted Index (like Elasticsearch or Lucene).

During search time, the query is passed through the same SPLADE model to expand the query terms, and the engine calculates the dot product using standard inverted index mechanics. You get the speed of BM25 with the semantic recall of dense embeddings.

Show Me the Code

While training a SPLADE model requires PyTorch and careful regularization, utilizing a pre-trained SPLADE model to generate sparse vectors is straightforward using the Hugging Face transformers library.

1import torch2from transformers import AutoModelForMaskedLM, AutoTokenizer3
4# Load a pre-trained SPLADE model5model_id = "naver/splade-cocondenser-ensembledistil"6tokenizer = AutoTokenizer.from_pretrained(model_id)7model = AutoModelForMaskedLM.from_pretrained(model_id)8
9document = "The automobile broke down on the highway."10
11# Tokenize the input12tokens = tokenizer(document, return_tensors="pt")13
14# Forward pass15with torch.no_grad():16    output = model(**tokens)17
18# Extract the logits (predictions for the vocabulary)19logits = output.logits20
21# SPLADE uses max pooling across the sequence length, followed by a ReLU22# to ensure all weights are positive and to induce sparsity23sparse_vector = torch.max(24    torch.log(1 + torch.relu(logits)) * tokens.attention_mask.unsqueeze(-1), 25    dim=126)[0].squeeze()27
28# Find the non-zero indices (the "expanded" words)29non_zero_indices = sparse_vector.nonzero().squeeze()30non_zero_weights = sparse_vector[non_zero_indices]31
32# Map the indices back to human-readable words33expanded_doc = {}34for idx, weight in zip(non_zero_indices, non_zero_weights):35    word = tokenizer.decode([idx])36    expanded_doc[word] = round(weight.item(), 3)37
38# Sort by weight to see the most important terms39sorted_doc = dict(sorted(expanded_doc.items(), key=lambda item: item[1], reverse=True))40
41# Print the top 10 terms42print(list(sorted_doc.items())[:10])43# Note how "car" is highly weighted despite not being in the original text!44# -> [('automobile', 2.87), ('car', 2.41), ('highway', 2.15), ('broke', 1.88), 45#     ('broken', 1.72), ('vehicle', 1.65), ('breakdown', 1.43), ('road', 1.21), 46#     ('down', 0.95), ('motor', 0.88)]

Watch Out For

Increased index size

While Learned Sparse Retrieval uses a fast inverted index, it is significantly heavier than traditional BM25. Because the model expands documents with dozens of synonyms and related terms, the inverted index must store far more entries. A SPLADE index can consume 3x to 5x more disk space and RAM than a standard BM25 index of the same corpus.

Query latency overhead

BM25 requires almost zero compute at query time (just string splitting and math). SPLADE requires running the user's query through a Transformer model (BERT) before the search can even begin. If you do not have GPU acceleration on your query servers, the neural network inference step will add 50-100ms of latency to every search.

The Quick Version

Lexical search (BM25) is fast but fails on synonyms. Dense embeddings understand synonyms but are computationally heavy and unexplainable.
Learned Sparse Retrieval (like SPLADE) uses a neural network to analyze a document and predict highly relevant synonyms.
Instead of outputting a dense vector, it outputs a list of weighted words (including words not in the original text).
This expanded list of words can be loaded into a standard, blazing-fast inverted index, granting semantic understanding to traditional keyword search infrastructure.