cat posts/2024-06-25-vector-embeddings-visually.md
Vector embeddings, visually
A support engineer typed “how do I cancel my subscription” into our new help center search last month, and the top result was the article on how to renew it. Second was the cancellation article.
That search runs on embeddings, the same kind that sit under the docs bot we built last fall and under most of the “AI search” features everyone has shipped in the last year. The pitch is always that embeddings find things that are similar in meaning, not just in spelling. That’s true, and it’s also where the trouble starts, because “similar” means something quite specific to the model, and it’s not always what a person means. So this is an attempt to draw what an embedding is, from a toy you can do on paper up to the models we actually call, and to pin down what “similar” means at each step.
A toy, by hand
The oldest idea here is from linguistics, not computer science. J. R. Firth’s line from 1957, “You shall know a word by the company it keeps”, is quoted in more or less every paper on the subject. Words that show up in similar contexts tend to mean similar things.
You can turn that straight into numbers. Take four words and count how often each one appears near four context words in some pile of text. These counts are made up but plausible:
| drink | hot | type | screen | |
|---|---|---|---|---|
| coffee | 8 | 6 | 1 | 0 |
| tea | 7 | 5 | 0 | 0 |
| keyboard | 0 | 0 | 9 | 2 |
| laptop | 1 | 1 | 6 | 8 |
Each row is now a vector, a point in four-dimensional space. That’s an embedding, a crude one. The question “are these two words similar?” becomes “do these two vectors point the same way?”, and the standard way to measure that is cosine similarity: the dot product divided by the two lengths, which is the cosine of the angle between them. (I walked through it by hand with three-dimensional vectors in the docs bot post, so here it’s words instead.)
Coffee and tea: the dot product is 8×7 + 6×5 + 1×0 + 0×0 = 86. The lengths are √101 ≈ 10.05 and √74 ≈ 8.60, so the cosine is 86 / 86.45 ≈ 0.99. Nearly the same direction.
Coffee and keyboard: the dot product is just 1×9 = 9, the keyboard’s length is √85 ≈ 9.22, and the cosine is 9 / 92.7 ≈ 0.10. Close to a right angle.
Keyboard and laptop come out at 0.75, related but not the same, and coffee and laptop at 0.20, because people do occasionally drink coffee near a laptop and the counts pick that up.
Notice that cosine ignores length. Tea’s vector is shorter than coffee’s because tea showed up less in my imaginary text, and it doesn’t matter. Only the direction counts, which is what you want, since a word being common shouldn’t make it more similar to everything.
The problem with counts is the size. With a real vocabulary you’d have one dimension per context word, tens of thousands of them, almost all zeros. Andrew Ng’s class in 2011 had a week on PCA, squashing a high-dimensional dataset down to a handful of directions that keep most of its variance, and if you do that to a big co-occurrence matrix you get something surprisingly good. That’s roughly what latent semantic analysis was doing in the ’90s.
word2vec, and directions that mean things
In 2013 Tomas Mikolov and others at Google published word2vec, which learns the vectors directly. You give each word a random vector of, say, 300 numbers, then train a small network to predict a word’s neighbors from the word (or the word from its neighbors), and nudge the vectors to make the predictions better. The vectors that fall out are dense, with every number doing some work, and they captured something the count tables didn’t make obvious: directions mean things.
The famous example is that vector(“king”) − vector(“man”) + vector(“woman”) lands closest to vector(“queen”). There’s a direction in the space that roughly corresponds to gender, another to country-and-capital, another to verb tense. The example is a bit cherry-picked (the published versions exclude the input words from the answer, otherwise “king” itself tends to win) and not every analogy works. But it changed how people thought about these things, from “a lookup table of numbers” to “a space with structure”.
It also had the problem my support engineer found, right from the start. Antonyms show up in nearly identical contexts. “The water is hot” and “the water is cold” are the same sentence with one word swapped, so “hot” and “cold” come out as close neighbors. By the company they keep, cancel and renew are practically the same word: they both appear next to “subscription”, “plan”, “account”, “billing”, “how do I”. A distributional model is learning about topic and role, and opposites share both.
Sentences, and what they’re trained to call similar
The embeddings we use at work aren’t word vectors. We send a whole chunk of text, a question or a paragraph, and get back one vector for all of it. OpenAI’s current models, text-embedding-3-small and text-embedding-3-large, came out in January and return 1,536 and 3,072 numbers. The small one costs two cents per million tokens, which is cheap enough that cost has basically stopped being a reason not to embed something.
These come from transformer models, and the important part is how they’re trained, because the training objective is the definition of “similar”. The usual method is contrastive. You collect a huge number of pairs that should be close, a question and the passage that answers it, a title and its article, two paraphrases of each other, and train the model so each pair’s vectors point the same way while unrelated texts in the same batch get pushed apart. After enough of that, “similar” means “the kind of thing that was paired together in training”, which for most of these models is roughly “about the same topic, or likely to answer the same question”.
Under that definition, the renew article is similar to “how do I cancel my subscription”. It’s about subscriptions, it’s a how-to, it would plausibly be in the same help center section, and the one word that flips the meaning is one word among many. Negation is a known weak spot for the same reason. “Articles that mention refunds” and “articles that don’t mention refunds” are very close vectors.
The annoying part is that we already knew the fix. On the docs bot last fall we’d learned that pure vector search falls down on exact words, and we’d ended up running BM25 alongside it. The help center search was built by a different team, from a vendor’s “semantic search” starter, and the hybrid never carried over, partly because nobody had written it down anywhere except in the bot’s code and on this blog. So the fix was the same one: a plain keyword search alongside the vector search, using the BM25 ranking I wrote about a couple of years ago, with the two result lists merged. The keyword side cares a lot about the literal word “cancel”, which is exactly what was missing, and it also handles the things embeddings are bad at, like error codes, product SKUs and the names of settings, where a near miss in meaning is a complete miss for the user. Embeddings handle the queries where the user’s words aren’t the article’s words.
A few things about high-dimensional space
This is the part where the math degree actually changes how I read the numbers.
In two dimensions, two random directions are as likely to be at 30 degrees as at 90. In 1,536 dimensions, that’s not true at all. Two random unit vectors there have a cosine that’s very close to zero, with a standard deviation of about 1/√1536 ≈ 0.026. Almost everything is almost perpendicular to almost everything else. So in a well-spread space, a cosine of 0.3 is a very strong signal, not a weak one.
Real embedding models aren’t well spread, though, which confuses people the other way. With OpenAI’s previous model, text-embedding-ada-002, nearly everything in our data scored above 0.7 against nearly everything else, because the vectors all sat in a fairly narrow cone. The ordering was still useful, but the absolute number wasn’t, and I’ve seen more than one team (including ours) set a “relevance threshold” of 0.8 and wonder why it let everything through. The lesson is to look at the distribution of scores for your own data before picking any cutoff, and to compare scores only against other scores from the same model.
The new OpenAI models also do something I found neat: you can ask for fewer dimensions, say 256 instead of 1,536, and they’re trained so that the first numbers carry the most information, a technique from a 2022 paper called Matryoshka representation learning. You lose a bit of accuracy and storage drops by six times. For our help center, with a few thousand articles, storage doesn’t matter. For anyone embedding a few hundred million documents it matters a lot.
Directions inside the model
One more thing, because it’s been the most interesting AI news of the last month for me. Anthropic published a paper in May, “Scaling Monosemanticity”, where they found millions of interpretable “features” inside Claude 3 Sonnet, directions in the model’s internal activations that correspond to specific concepts, by training a sparse autoencoder on them. One responds to the Golden Gate Bridge, and for a short while they let people try a version of Claude with that feature turned up, which talked about the bridge no matter what you asked it. It’s not quite the same thing as an embedding model, but it’s the same family of idea as king − man + woman, at enormous scale: concepts are directions, and you can find them, and push on them.
Which means the thing my support engineer screenshotted isn’t a bug in the usual sense. The model put “cancel” and “renew” close together because, along most of the directions it knows about, they are close. The direction they differ on just isn’t one it weighs heavily for search, and the only way I know to fix that from the outside is to give it some help from something that reads the words.