DEV Community

Cover image for An embedding is a lookup table, and everything else is how you fill it
Carlos Chinchilla Corbacho
Carlos Chinchilla Corbacho

Posted on Edited on Originally published at cchinchilla.dev

An embedding is a lookup table, and everything else is how you fill it

Originally published at cchinchilla.dev. Part 2 of From code to weights, a 12-part series on ML fundamentals for engineers.

Part 1 ended on an index. The tokenizer turns text into ids, an id is a row number, and the row lives in a matrix the model was born with. This post reads that matrix. It's smaller than it sounds: a lookup table, indexed by integer, with learned contents. The science is in how the contents get there. The engineering is in what you can't do with the table once they're in.

I read two files for this post: the forward of torch.nn.Embedding, which is one call, and the pooling module of sentence-transformers, which is where the other kind of embedding, the one you put in an index, gets made. Then I built a 48-line index and broke it three ways on purpose, to measure what breaks.

This is part of a series for engineers moving into ML engineering. Each part names one thing that does not behave the way software behaves. This one is the index: it has no migration.

A matrix and an index

If you've written a hash map, you've written most of an embedding layer. nn.Embedding(V, d) allocates a matrix of V rows and d columns. Looking up id i returns row i. That's the whole forward pass, and torch's own source says so in sparse.py:188-197:

def forward(self, input: Tensor) -> Tensor:
    return F.embedding(
        input,
        self.weight,
        self.padding_idx,
        self.max_norm,
        self.norm_type,
        self.scale_grad_by_freq,
        self.sparse,
    )
Enter fullscreen mode Exit fullscreen mode

F.embedding is a gather. That's checkable rather than a matter of trust: for a fresh Embedding(50257, 768) and three ids, torch.equal(emb(ids), emb.weight[ids]) is True. No transformation, no hashing, no similarity. A row per id.

The sizes are where it stops looking like a hash map. GPT-2's input table, wte.weight, is 50,257 rows by 768 columns: 38.6 million numbers, 154 MB in fp32, every one of them a trained parameter. all-MiniLM-L6-v2, the sentence model the rest of this post runs on, has a 30,522 by 384 table for the same job. Row 11,241 of GPT-2's table is the vector for " token" for the reason Part 1 gave: that's where the tokenizer sends it. The contents of the row were learned. The address was fixed before training started.

For the engineer coming from software, the surprise is that the keys mean nothing. A hash map's key carries its own identity; row 11,241 means something only inside the model that trained it. For the data scientist, the surprise runs the other way. This table isn't a word2vec artifact you produce and then use. It's a parameter inside the model, updated by the same gradient step as everything above it. It never leaves.

Two ways to fill the table

A row gets its numbers in one of two ways, and the difference decides what you can do with the result.

The first is the one above. The table is a parameter of a language model. Rows start random and gradients flow into them through the layers above. After training, each row is whatever helped the model predict the next token. Nobody designs the space; it's a side effect. Part 8 reads the gradient that fills it.

The second is the one most engineers actually deploy. You take a trained model, run a piece of text through it, and take the output as the text's vector. The row isn't stored anywhere; it's computed on demand, and it depends on the whole input, not on one id.

The word bank has one row in the input table of that sentence model, row 2,924. A position table gets added to it, but the word's row is the same in "she sat on the river bank" and in "she opened a bank account". At the output, the two vectors for that same token have a cosine of 0.77 to each other. Same lookup, different result, because a position and six layers sit between the two. That's what contextual means, and it's the whole difference between a static embedding and the kind you put in a search index.

The output is a vector per token. A sentence needs one vector, so something has to collapse T vectors, one per token, into one, and that something is pooling. sentence-transformers keeps every variant in one place, and on the default padded path the branch a mean-pooled model runs is pooling.py:200-213:

elif mode in ("mean", "mean_sqrt_len_tokens"):
    if mean_sum is None:
        mask = attention_mask.unsqueeze(-1).expand_as(token_embeddings).to(token_embeddings.dtype)
        mean_sum = (token_embeddings * mask).sum(dim=1)
        # ... mean_mask is the count of real tokens (or a supplied per-token weight sum), clamped away from zero
    if mode == "mean":
        output_vectors.append(mean_sum / mean_mask)
Enter fullscreen mode Exit fullscreen mode

Sum the tokens the attention mask says are real, divide by how many there were. Six modes live in that function, each a different way to collapse the same T vectors: cls takes the first, lasttoken the last, max the largest per dimension. A model is trained with one of them, and the choice is part of the artifact.

So is the step after it. The small model's own modules.json lists what runs when you call encode:

[
  {"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"},
  {"idx": 1, "name": "1", "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"},
  {"idx": 2, "name": "2", "path": "2_Normalize", "type": "sentence_transformers.models.Normalize"}
]
Enter fullscreen mode Exit fullscreen mode

Transformer, then pooling, then normalisation to unit length. Not even the family agrees on the third: paraphrase-MiniLM-L6-v2 ships without it, and its raw vectors have norms around six. Three modules, and the last two are the ones a client tends to reimplement. Reimplemented differently, they change the answer. By how much is measured two sections down.

Close is a dot product

Once every vector is unit length, cosine similarity is a dot product, and nearest-neighbour search is a matrix multiply followed by a sort. GPT-2 reads its own table the same way on the way out: the logits for the next token are a dot product against every row of wte. That's the whole retrieval side, and it's why the index bill is arithmetic: ten million documents at 384 dimensions is 15.4 GB of fp32 before any index structure. At 768 it's 30.7 GB. At 1,536, the width of a popular hosted embedding, 61.4 GB. Dimension is a cost you pay per document for as long as the index exists. Whoever trained the model decided it, not you.

The number the search returns is less useful than it looks. On my 48-line corpus, the best match for "how do I undo a bad release" has a cosine of 0.30 to the query, and it's the right line: the rollback procedure. The best match for "the migration is stuck waiting on a lock" scores 0.71. Both are correct. A cosine isn't a confidence; it's a position in a space whose scale the model chose. Compare ranks within one model, and never compare the raw number across two.

There is no schema migration

Here is the part that's different from software. When you change a database schema you write a migration, run it, and the old rows come out the other side in the new shape. An embedding space has no migration. Two models produce two spaces, and no map between them ships with either, because neither was designed; each is whatever training settled on. You can fit one after the fact, and people do. It's a learned approximation with its own error, not a migration: a linear map fitted on four thousand unrelated sentences lifts the stale index below from 60% to 70%, and stops there. Re-embed with a new model and every vector you stored is in a coordinate system nothing else uses.

I expected that to show up as garbage. It shows up as something worse.

The test: 48 lines from a runbook, eight each on deploys, auth, billing, databases, observability, and on-call. I indexed them with the L6 model and queried through all-MiniLM-L12-v2, the upgrade an engineer would reach for (same family, same 384 dimensions, so the index accepts the vectors without complaint).

Against what the upgraded index would return, the stale one keeps 60% of the top five. Random would keep 10%, and guessing five lines from the right topic would keep 35%. The old index queried by the old model keeps 77%: the mismatch costs seventeen points, and the other twenty-three are the two models disagreeing about the answer. For "user keeps getting logged out", the first hit is still the session-expiry line, and the third is the database backup retention policy. Nothing in the response says which is which.

Be fair about the number. Sixty percent isn't a scrambled index, and the reason is lineage, not English: both models were fine-tuned from the same parent checkpoint, the L6 one after dropping every second layer, so the two spaces started in the same coordinates. The same line under the two models has, on average, a cosine of 0.54 to itself, far from orthogonal. Two models that only shared training data wouldn't share coordinates, and a stale index between them should land near random. Six tight topics of eight lines each is also the friendliest corpus a stale index could hope for. I haven't measured a production index, and I'd expect it to keep less. What the number does support is the shape: partial, plausible, silent.

The pooling module does the same thing without a model change. Take the same weights and pool with cls instead of mean, at index time and at query time, and the top five keeps 50% of the mean-pooled answer. Mix them, mean in the index and cls on the query path, and it keeps 70%: more than consistent cls, because one side of the pair is still the pooling the model was trained with. A client that reimplements pooling has changed the model.

There's an interactive version of the next bit that doesn't survive being pasted into a feed, so it lives on the original post: drag the reindex forward, then flip which model encodes the queries. On the new model, recall climbs as the table migrates. On the old one it starts at 0.77, because two consistent models already disagree on a quarter of the answers, then falls through the middle, to 0.63, gets back to 0.77 only in the last few lines, and finishes below it, at 0.73: by the end, every line in the index was written by a model the query path isn't using.

That second curve is the one production runs. The index is a batch job; the query encoder is a deployed service; the two ship on different days.

What the ML engineer watches

Four habits, all cheap next to a reindex.

The embedder is part of the index. Store the model's hash, its pooling mode, its normalisation step, and its sequence cap in the index metadata. Refuse a query whose encoder doesn't match. The modules.json above is half of that contract; the pooling mode sits next to it in 1_Pooling/config.json, and the hash you add yourself. The cap matters: the L6 model reads 256 tokens and the L12 upgrade reads 128, so the swap that broke the index also halved what each vector sees, silently. Ship all of it with the index, not in a client's config.

Cut over, don't roll. Build the new index beside the old one, switch queries and index together behind one flag, and delete the old one afterwards. A rolling reindex under a live query path is the falling curve.

Measure neighbour overlap, not cosine. Keep a fixed set of queries with their expected top-k, and check recall@k against it on every model change. A cosine threshold tuned on one model means nothing on the next; the neighbour set is the only observable that survives a swap. It measures drift, not relevance: the migrated top five for the logout query still has a cold-storage line in it.

Dimension is a bill. N times d times four bytes, before the index structure. Half-precision halves it; some newer models let you keep only the first k dimensions. Both can be applied to vectors already stored, without re-embedding; truncated vectors have to be renormalised and queries cut to the same k, and both choices belong in the index metadata with the rest. Price a larger model's dimension against the corpus you'll have in two years, not the one you have now.

None of this is modelling work either. An embedding index looks like a database with a float column, right up to the day someone upgrades the model that wrote the column.

Part 2 of twelve. An embedding is a row of a matrix, filled either by the gradient or by a model's output, and the space it lives in has no migration: change the model and every stored vector goes stale, silently. A batch of sequences of those rows is a tensor with the signature [B, T, d]. The next part reads that signature the way you'd read a function's, and finds that the runtime never checks it.

Written with AI assistance. See AI policy.

Top comments (4)

Collapse
 
mihai_leanzero profile image
Mihai Perdum •

Carlos, the embedder-as-part-of-the-index point is the one people skip until it bites. Your fix works cleanly for a local sentence-transformers checkpoint where pooling mode and normalisation are sitting right there in modules.json -- what happens against a hosted embedding API (OpenAI, Cohere) where none of that is published? You can still hash the returned vectors' dimension and a fixed probe set, but you can't introspect pooling or normalisation to put in the metadata the way you can locally, so "refuse a query whose encoder doesn't match" has less to check against. The 60% recall on a same-family stale index is lower than I'd have guessed given a 0.54 self-cosine across the two models.

Collapse
 
cchinchilladev profile image
Carlos Chinchilla Corbacho •

The objection is right, and it changes the instrument rather than the habit. You can't read pooling or normalisation off a hosted API, so stop describing the encoder and fingerprint it instead.

Keep a probe set: thirty-two sentences from the corpus, not generic ones. Embed them when the index is built and store the vectors in the index metadata. The query service embeds the same probes at startup and on a schedule, and compares per-probe cosine against what is stored. Compare with a tolerance rather than for equality — batching and floating-point order move the last digits — so > 0.9999 per probe. The whole check is one API call.

That is a stronger contract than the local one, not a weaker substitute for it. modules.json reports the configuration; it says nothing about the weights. A provider that reships text-embedding-3-small under the same name is invisible to a config hash and loud to a probe set. Locally the two belong together: the config catches the client that reimplements pooling, the probes catch everything else.

Two parameters deserve the same treatment, because they are the hosted form of the pooling mistake. Dimension truncation on text-embedding-3-* is part of the artifact, not a client setting. And Cohere's input_type is asymmetric by design: search_document when you index, search_query when you query. Swap them and you have built mean-in-the-index, cls-on-the-query, inside a wrapper you didn't write.

On the 60%: the ceiling is 77%, not 100%. That is what the old index keeps when the old model queries it, because two consistent models already disagree on a quarter of the answers. The mismatch costs seventeen points.

Self-cosine doesn't predict recall, and the reason is that recall@5 is decided by ordering. 0.54 says how far a line moved. Whether it stays in the top five says whether it moved differently from its competitors. Six topics of eight lines puts the boundary among near-ties, which is where a small inconsistent perturbation flips ranks; a large consistent displacement would leave recall untouched.

Be fair about the corpus. The L12 upgrade also halves the sequence cap, 256 to 128, and my runbook lines are short enough that it never bites. On real documents it would, and the response would not mention it.

Collapse
 
mihai_leanzero profile image
Mihai Perdum •

This is a much better answer than what I was picturing, thanks for laying it out. Fingerprinting the encoder with a stored probe set instead of trying to read its config off an API is obviously the right move once you say it - you can't inspect what you can't see, so stop trying to and just measure the thing directly.

The dimension-truncation and input_type points are ones I would have gotten wrong on my own - I was thinking of hosted embedding APIs as one flavor of "opaque," and you're right that they carry their own footguns beyond the pooling/normalisation gap I was asking about, not just weaker versions of the local one.

One thing I'm still turning over: the probe set itself has to be representative of the corpus to mean anything, per what you said. What happens when the corpus drifts enough that the original 32 sentences stop being representative queries - does the tolerance check degrade quietly the same way the underlying recall does, or is there a separate signal for "the probes are stale now too"?

Thread Thread
 
cchinchilladev profile image
Carlos Chinchilla Corbacho •

No, and that's worth separating, because the two failures aren't the same instrument.

The probe check measures the model, not the corpus. Same weights, same probes, cosine sits at 1.0 whether or not those thirty-two sentences still look anything like what you index today. Corpus drift can't make it quietly wrong.

What drift does cost you is detection power. Probes drawn from the corpus aren't there for correctness, they're there for sensitivity: a fine-tune, a quantisation, or a silent version bump on a hosted endpoint can be indistinguishable from the original on generic text and clearly different on your vocabulary. Stale probes stop exercising the part of the space where a near-identical checkpoint gives itself away, so the check keeps passing when it should be the thing that catches it. The failure mode isn't a false alarm, it's a false pass — which is the silent degradation you're asking about, one level up.

So yes, a separate signal, and not a calendar. Two sets with different jobs:

A rolling set, resampled from the corpus on every index build and stored with it. The probes are then never older than the index they guard, which is the case that matters, because the probe set is an artefact of the index.

A small frozen anchor set that never changes, for identity across rebuilds. It's the only thing that can tell you the encoder moved between two indexes rather than within one.

And if you want the staleness measured rather than assumed: take the live query log, embed a sample, and watch the fraction whose nearest probe falls below some cosine. When that fraction climbs, the rolling set has stopped covering what people actually ask, and you learn it from traffic instead of from a rebuild schedule.