🧵 Stella Embeddings: What's the big deal? (a mini explanation thread)

If you enjoy RAG Twitter, or compulsively check the MTEB leaderboard, you might've come across the "Stella" (and now, Jasper) models.

They seem suspiciously good, even though you've never heard of them before. Well, let's begin with a spoiler: they're good, period.

✨ Stella is really cool, for a few reasons:
- It's a piece of independent research into topics that aren't the current mainstream way of doing things, with fantastic results to back it up
- It's specifically designed to avoid overfitting and maximize generalization

🏆 Let's take a step back to understand a big piece of signal that makes the Stella embeddings special: MTEB

MTEB is the Massive Text Embedding Benchmark. It's a set of curated datasets designed to evaluate embedding models' capabilities on various tasks, including a mix of retrieval tasks. It's essentially the go-to when you're trying to figure out what embedding model you should try.

🤔 If you occasionally glance at the IR people Twitter discussions, you'll see that MTEB is a pretty disputed hot topic. That's because there are a lot of different setups co-existing on a single leaderboard:
- Some models attempt to be truly zero-shot (MTEB was designed as a zero-shot eval, but has a sizeable trainset)
- Some models train on some subsets/all of it, in various ways.

⚠️ Don't get me wrong: training on the training set is good, actually. It's valuable signal towards the potential of a model when finetuned on domain-appropriate data. There's absolutely nothing wrong with overfitting MTEB, in a vacuum.

The problem is that they're all blended together on the same leaderboard: it's impossible to tell who's doing what, and what level of contamination you should expect for each model. Is this model lagging behind because it's worse? Or because it's fully zero-shot? Nobody knows.

🔍 Now then, let's look at how embedding models are usually trained:

The current dominant way of training retrieval models is via the use of a contrastive loss, with little-to-no knowledge distillation (except some, such as ColBERT and ColBERTv2, or RocketQA, but it's very rarely used by MTEB mainstays).

The contrastive process uses Query <-> Document relationships, with a two-step approach:
1️⃣ First, use a lot of distantly supervised data to align the model's embedding space onto the concept of "similarity" between documents and queries. This'll be stuff like a Reddit title and its comments, Title <-> Abstract pairs from SemanticScholars, etc.
2️⃣ Then, fine-tune on higher quality annotated data, either extracted from a search engine based on which documents users clicked on after their search. Or, more generally, any dataset where human annotators labeled that the document is an answer to the query.

💭 In my opinion (and I backed it up with answerai-colbert-small!) distillation should have a bigger place. This is especially true nowadays that we have LLM-based embedders, which can be very good, but also very, very impractical given their size.

✨ And you know what? This is where Stella enters the arena!

Stella's approach takes distillation even further than previous methods: what if we only use unsupervised distillation? Why don't we get rid of those expensive, pesky labels, and complicated hard negative mining pipelines? (spoilers: this is a good idea):

🔮 Essentially, Stella takes a few really good 7B LLM-based embedding models and comes up with a multi-step learning process, involving freezing/unfreezing, dimension manipulation tricks, and much more: read the paper to find out!

📐 The main aim of this training is to work within the embedding space, seeking to minimize the geometric distances (both cosine and angular distance) between the teachers' vectors and the student model (Stella)'s outputs.

🤖 The resulting model is one that has never itself been trained on predicting query <-> document relationships like a normal retrieval model would, but instead has learned to mimic how models much smarter than itself would represent both queries and documents.

🎯 In practice, Stella models generalize amazingly well because of this. Most other small models overfit on MTEB in some ways, consciously or not: either by training on its data or, in the case of some models, by iterating on the synthetic data mix they use and selecting the one that has the best MTEB dev score.
This generally makes it kinda hard to know what's the best model for downstream use: is the model really good at generalization (what 99% of people want for their own use cases) or is the model really well fitted to MTEB? (As discussed above, this is useful signal in itself, but a lot of people want to use off-the-shelf embeddings and that's it, in which case, this is not an appropriate signal)

💫 Stella doesn't really have that problem: it's good at MTEB because it's good at generating representations that look like its teacher models' representations, and its teacher models' representations are good because they are gigantic models that are at least 7x, oftentimes 20-30x, bigger than common retrieval models.

Stella isn't trained on MTEB; at worst, it's picked up some of what its teachers knew about it. But thankfully, the teachers selected weren't overfitted on MTEB.

🎉 This means that the signal you get from Stella doing well on MTEB is that much more meaningful: you know for a fact it's good at generalizing, and so intuitively, it's a good starting point for your task.

(This has been anecdotally verified quite a few times -- on the new TREC RAG datasets, and I've heard from a bunch of people in various industries who got great results by just swapping in the latest Stella model instead of whatever they were using before)