Exploring Static Embedding Retrieval

At LlamaIndex, we are huge fans of static embeddings. We've used them to power past projects like SemTools, various features in the LlamaParse platform, as well as heavily iterating with them for other upcoming features.

Static embedding models can seem like a cheat code at first glance. Embedding a line of text is just a table lookup plus an average calculation. This translates to about 0.05 ms per text line on CPU, no transformer forward pass, it's WASM-friendly, and ~100× faster than even a small dense model.

We spent a couple of weeks trying to get late-interaction retrieval (ColBERT-style MaxSim scoring) out of a static embedding model, and basically figured out it doesn’t work in several ways. But we wanted to share our exploration anyways, because we think static embeddings are just that cool.

What are static embeddings?

Static embeddings have a long history, going from being one of the earliest ways to embed text, to being the fastest way to embed text today. Modern approaches work by generating an embedding for every token in a models vocabulary, and saving that to a look-up table (i.e. Model2Vec distillation). Approaches like tokenlearn then optimize these by embedding chunks with the parent model and moving the average static embedding for that chunk to more closely match the parent model.

However, static embeddings have two main drawbacks:

1. Embeddings are only per-token, not contextual. This means phrases like "not good food" have high similarity with "good food"
2. Embeddings are pooled and averaged. This means longer text lines have diluted meaning. To be fair, typical dense embeddings have a similar issue, but this is amplified by issue #1.

ColBERT models and others like it use a process called late-interaction, and can avoid the pooling problem by keeping one vector per token and scoring with MaxSim. This is a scoring process where each query token finds its best-matching document token, and those best-match scores are summed.

Our thinking was, a static model already has per-token vectors in its lookup table. Could we skip pooling, score with MaxSim, and get late-interaction accuracy? What would it take to get there? We took the excellent minishlab/potion-retrieval-32M model as our base for the following attempts.

Attempt #1: MaxSim over raw static tokens

The first experiment to skip the pooling was simply to keep the per-token static vectors, and score with MaxSim directly.

The results were mixed:

• On a small internal set of hard exact-match queries where the answer is a specific string buried in a document, and the query shares tokens with it. MaxSim showed great improvement over pooling (R@3 0.38 vs 0.14). The token averaging that pooling requires was destroying the per-token signal that MaxSim naturally captures.
• On paraphrased queries MaxSim was no better than pooling, and ranking quality actually got worse. And on public benchmarks it lost to pooling (0.418 vs 0.504 mean NDCG, full table below).

In hindsight this makes sense. MaxSim works in ColBERT because a transformer has already adjusted each token's vector based on its context. Static vectors have no context, which means every occurrence of a token gets the same vector no matter what surrounds it.

Attempt #2: teach the tokens some context

So the vectors need context. In the spirit of keeping things light and fast, our idea was to add a tiny "mixer" that runs over a line's token vectors after lookup and before MaxSim, nudging each token's vector based on its neighbours. As a mixer, we trained a small convolution head:

# tokens: the line's static vectors, one 512-dim row per token
# each token blends with 3 neighbours on each side
mixed = depthwise_conv(tokens, kernel_size=7)

# mix information across the 512 dimensions
mixed = linear(mixed)

# per-token: how much mixing to accept                       
gate  = sigmoid(linear(tokens))

# small nudge on top of the original vector                
out   = l2_normalize(tokens + gate * mixed)    

This mixer is ~530k parameters or ~2 MB. Inference is a couple of small matrix multiplies per token which was microseconds per line, and it did not impact the speed that made static embeddings attractive in the first place.

We trained it by distillation: run C4 text through a tradition dense embedding model (bge-base-en-v1.5), and taught the mixer to nudge each static token vector toward the direction of the transformer output.

It was working? Kinda?

On NanoBEIR (13 retrieval benchmarks, all three systems under one harness, NDCG@10):

We can see that the conv mixer is providing some improvements. A score of 0.526 is 94% of MiniLM-L6 (0.562) at roughly 100× the speed, from a 2 MB add-on. But this is still far behind modern-ish dense models and late-interaction models. At this point we thought we just needed to train it better, and tried a few more ideas.

Attempt #3: a smarter teacher

We were distilling from bge-base, a model whose per-token vectors were never trained to be used for MaxSim (only its pooled output is optimized). ColBERTv2's token vectors are trained for MaxSim and so swapping the teacher seemed like the right call.

The ColBERT-taught mixer improved its validation accuracy (val cosine 0.615 vs 0.32), but then retrieved worse overall (0.504, tying plain pooling).

ColBERT's vectors get their power from deep, context-specific structure that a small conv over frozen vectors can't reproduce. It can match the average direction and still miss everything that matters when doing MaxSim retrieval later on.

Attempt #4: skip the teacher, train on the actual metric

We trained the same mixer directly on retrieval: a contrastive loss over MaxSim scores on MS MARCO query/passage pairs.

Result: 0.5267 while the original distilled conv mixer was 0.5258., essentially a tie.

Attempt #5: fine-tune the table itself

So far, the original static-embedding table has been frozen while we train additional parameters from the conv layer. If the frozen table is the bottleneck, train the table. We fine-tuned the full 32M-parameter embedding table on the same contrastive MaxSim objective.

Two bad things happened. MaxSim barely improved (0.418 → 0.435, nowhere near the conv's 0.527). And pooled quality dropped (0.504 → 0.487). We realized that MaxSim only routes gradient to a handful of winning tokens per example, and spreading that sparse signal across 32 million rows goes nowhere during the training process.

Takeaways for future exploration

Even if this exploration didn’t produce results that we liked, it was still a fun dive into exploring different ways to push static embeddings to their limits. A few lessons we learned:

• Don't score raw static token vectors with MaxSim. It's worse than the standard pooling, the scoring method is technically changing the training target. MaxSim assumes contextualized vectors and static embeddings are not contextual.
• A tiny conv adapter can technically add semantic context . But it also doesn’t compare to standard transformers and attention. We probably could have spent more time improving this adapter layer, but losing out on model swapping didn’t seem worth it.
• A better teacher is not a better student. When the student is too small to absorb what the teacher knows (i.e. a small conv model), teacher quality stops mattering.

In our view, improvements to static embeddings are not going to come from these approaches. But by sharing these experiments, hopefully we can inspire future research directions for scalable and lightning fast embeddings.