Embeddings


< all posts

Published : 2026-07-11

Embeddings

Suppose we have some data which is in text form, now models can't perform any arithmetic on text so we need to convert it into a representation which enables a model to actually parse the info.

Let's take an example of a data which contains three cities :

cities = [New York, Beijing, Tokyo]

in order to assign some numeric structure we can assign numbers to these cities :

cities = { New York : 1, Beijing : 2, Tokyo : 3 }

Assigning New York = 1, Beijing = 2, Tokyo = 3 secretly tells the model three lies: that Tokyo (3) is somehow "greater than" New York (1), that Beijing is exactly "between" them, and that Tokyo = 3 × New York. Any model that does arithmetic on its inputs will act on this fake structure. The numbers imply an ordering and a scale the cities don't actually have.

One Hot encoding

For categorcial variable with KK values, map category kk to the standard basis vector :

ek{0,1}Ke_k \in \{0,1\}^K

So our entry ekje_{kj} returns 1 if j==kj == k and 0 otherwise. Cities defined above now become :

cities = { 
    New York :  (0,0,1), 
    Beijing :   (0,1,0), 
    Tokyo :     (1,0,0) 
}

In this type of encoding every category is equally far from one another and this encoding asserts that no two categories are related to each other. This is good but not the best way as if we have K=50000K = 50000 then it is not the best way. We can actually make a learned vector representation.

Embedding as a learned geometry

So we want the best of both worlds: a representation that's compact (not 50,000 dimensions) and can express similarity (not everything orthogonal). That is exactly what an embedding is.

Instead of a fixed 0/1 vector, we give every category its own short, dense vector of length dd, and we collect all of them into a single matrix named embedding table:

ERK×d,dKE \in \mathbb{R}^{K \times d}, \qquad d \ll K

One row per category. Row kk is the vector for category kk:

emb(k)=EkRd\text{emb}(k) = E_k \in \mathbb{R}^d

For our three cities with, say, d=2d = 2, the table might look like:

E =  New York : ( 0.9,  0.2)
     Beijing  : (-0.7,  0.8)
     Tokyo    : (-0.6,  0.7)

Notice what this buys us that one-hot couldn't: Beijing and Tokyo have nearby vectors, while New York sits off on its own. The representation can now say "these two are similar" which is something one-hot's forced orthogonality made impossible. And dd is our choice (here 2, in practice often 100–300), completely decoupled from how many categories KK there are.

An embedding is just a one-hot times a matrix

Grabbing "row kk" looks like plain array indexing, but it's exactly the matrix multiplication we already met in §1. Take the one-hot vector ek\mathbf{e}_k (all zeros except a 1 at position kk) and multiply:

emb(k)=Eek\text{emb}(k) = E^\top \mathbf{e}_k

The one-hot acts as a selector: multiplying picks out row kk and zeroes everything else.

Why Eᵀeₖ selects row k

Look at component jj of the result:

(Eek)j=m=1KEmj(ek)m.(E^\top \mathbf{e}_k)_j = \sum_{m=1}^{K} E_{mj}\,(\mathbf{e}_k)_m.

But ek\mathbf{e}_k is zero everywhere except position kk, where it is 1. So every term in the sum dies except the m=km = k term:

(Eek)j=Ekj.(E^\top \mathbf{e}_k)_j = E_{kj}.

That holds for every component jj, so the whole result is row kk of EE which is exactly emb(k)\text{emb}(k).

There's no new machinery. It's the same EekE^\top \mathbf{e}_k, The only difference is implementation: since the input is one-hot, we skip the multiply-by-zeros and just fetch the row directly. That turns an O(Kd)O(Kd) matrix–vector product (almost all of it multiplying by zero) into an O(d)O(d) row lookup. Same math, far cheaper.

And because it's a linear layer, it's trainable which mean that the entries of EE are ordinary parameters, and gradient descent can adjust them just like any other weights. That's what makes the vectors learned rather than hand-designed.

The size win, in numbers

This also settles one-hot's first failure. Suppose K=50,000K = 50{,}000 and we feed the representation into a hidden layer of size h=512h = 512.

Similarity is now something we can measure

The deeper win is one-hot's second failure which no shared structure. With embeddings, "how related are two categories?" becomes an actual computation on their vectors, the cosine similarity:

sim(k,l)=Ek,ElEkEl\text{sim}(k, l) = \frac{\langle E_k, E_l \rangle}{\|E_k\|\,\|E_l\|}

This is 1 when two vectors point the same way, 0 when orthogonal, 1-1 when opposite. Unlike one-hot where every distinct pair scored 0, "cat" as far from "dog" as from "Tuesday" the geometry here can place related categories close together and unrelated ones far apart.

How Embeddings Learn

The idea: a word is known by the company it keeps

This is the distributional hypothesis: words that appear in similar contexts tend to mean similar things. Coffee and tea both show up near "drink," "hot," "cup," "morning." If we could force a word's vector to be good at predicting the words around it, then words with similar neighbourhoods would be pushed toward similar vectors, without ever labelling anything as "similar."

Skip-gram turns that into a concrete prediction task: given a center word, predict the words around it. Slide a window over a corpus and read off (center, context) pairs. From "the queen ruled her kingdom," with center word ruled:

(ruled, the)   (ruled, queen)   (ruled, her)   (ruled, kingdom)

Do this across billions of words and you have a mountain of training pairs, each one a tiny nudge on the vectors.

The cast of symbols

Before the model, let's fix every symbol

SymbolShapeMeaning
VVscalarvocabulary size (this is the KK from §1–§2 — now, specifically, the number of distinct words)
ddscalarembedding dimension
ccindexposition of the center word in the vocabulary
ooindexposition of the actual observed context word (our target)
wwindexa running index over the whole vocabulary (the summation variable)
EERV×d\mathbb{R}^{V \times d}the center embedding table from §2
UURV×d\mathbb{R}^{V \times d}a second, context embedding table
vc=Ec\mathbf{v}_c = E_cRd\mathbb{R}^{d}the center word's embedding (row cc of EE)
uw=Uw\mathbf{u}_w = U_wRd\mathbb{R}^{d}word ww's context embedding (row ww of UU)

Why two tables? Because a word plays two roles:

and there's no reason those two roles should share the same geometry. So each word owns two vectors: one row in EE (as a center) and one row in UU (as a context). It's a modeling choice, and it makes the math clean.

The model: score, then normalize

We want, for a center word cc, a probability for every possible context word. The natural compatibility score between two vectors is their dot product, so define the score of word ww as

zw=uwvcz_w = \mathbf{u}_w^\top \mathbf{v}_c

large when uw\mathbf{u}_w and vc\mathbf{v}_c point the same way. Stacking this over all VV words is just the matrix–vector product z=Uvc\mathbf{z} = U\mathbf{v}_c. Then turn the scores into a probability distribution with softmax:

p(oc)=exp(uovc)w=1Vexp(uwvc)p(o \mid c) = \frac{\exp(\mathbf{u}_o^\top \mathbf{v}_c)}{\sum_{w=1}^{V} \exp(\mathbf{u}_w^\top \mathbf{v}_c)}

Exponentiating forces every score positive; dividing by the sum forces the total to 1. This is exactly softmax regression except the "features" (vc\mathbf{v}_c) are themselves being learned.

The loss: maximum likelihood, again

We want the model to assign high probability to the context word that actually occurred. That's maximum likelihood, and exactly as in the Loss and Gradients article taking the negative log turns it into a loss to minimize:

=logp(oc)\ell = -\log p(o \mid c)

This is the same negative-log-likelihood recipe: choosing this loss is choosing a categorical model of the data.

The gradient: a force field

To train vc\mathbf{v}_c we need vc\nabla_{\mathbf{v}_c}\,\ell. First rewrite the loss into a shape that's easy to differentiate:

substitute the softmax and simplify with logAB=logAlogB\log\frac{A}{B} = \log A - \log B and logex=x\log e^x = x:

=uovc+logw=1Vexp(uwvc)\ell = -\mathbf{u}_o^\top \mathbf{v}_c + \log \sum_{w=1}^{V} \exp(\mathbf{u}_w^\top \mathbf{v}_c)

Two clean pieces: a linear term, and a "log-sum-exp" term. Differentiating gives

  vc=uo+w=1Vp(wc)uw=Ewp(c)[uw]what the model expectsuowhat actually happened  \boxed{\;\nabla_{\mathbf{v}_c}\,\ell = -\mathbf{u}_o + \sum_{w=1}^{V} p(w \mid c)\,\mathbf{u}_w = \underbrace{\mathbb{E}_{w \sim p(\cdot \mid c)}[\mathbf{u}_w]}_{\text{what the model expects}} - \underbrace{\mathbf{u}_o}_{\text{what actually happened}}\;}
Deriving the gradient (the log-sum-exp step)

The first term is linear in vc\mathbf{v}_c, so by the rule v(bv)=b\nabla_{\mathbf{v}}(\mathbf{b}^\top\mathbf{v}) = \mathbf{b} its gradient is just uo-\mathbf{u}_o.

For the second term, apply the chain rule through the log\log and the exp\exp:

vclog ⁣weuwvc=weuwvcuwweuwvc=weuwvcweuwvcuw=wp(wc)uw.\nabla_{\mathbf{v}_c} \log\!\sum_w e^{\mathbf{u}_w^\top \mathbf{v}_c} = \frac{\sum_w e^{\mathbf{u}_w^\top \mathbf{v}_c}\,\mathbf{u}_w}{\sum_{w'} e^{\mathbf{u}_{w'}^\top \mathbf{v}_c}} = \sum_w \frac{e^{\mathbf{u}_w^\top \mathbf{v}_c}}{\sum_{w'} e^{\mathbf{u}_{w'}^\top \mathbf{v}_c}}\,\mathbf{u}_w = \sum_w p(w \mid c)\,\mathbf{u}_w.

The fraction is exactly the softmax p(wc)p(w \mid c), so the gradient of the log-sum-exp term is the probability-weighted average of the context vectors. Combining the two terms gives the boxed result.

Read the boxed gradient as a physical force. Gradient descent steps against it:

vcvcηvc=vc+ηuoη ⁣wp(wc)uw\mathbf{v}_c \leftarrow \mathbf{v}_c - \eta\,\nabla_{\mathbf{v}_c}\ell = \mathbf{v}_c + \eta\,\mathbf{u}_o - \eta\!\sum_w p(w \mid c)\,\mathbf{u}_w

Now play this over billions of pairs. Queen and king both co-occur with {royal, throne, crown, reign}, so both get pulled toward the same set of context vectors and vectors pulled toward the same targets drift toward each other. Similarity emerges purely from shared company; we never once told the model that two words are related.

The same "prediction − target" pattern. Group the gradient by word instead: vc=w(p(wc)1[w=o])uw=U(p^eo)\nabla_{\mathbf{v}_c}\ell = \sum_w \big(p(w\mid c) - \mathbb{1}[w=o]\big)\mathbf{u}_w = U^\top(\hat{\mathbf{p}} - \mathbf{e}_o). The coefficient on each context vector is predicted probability minus target — the exact softmax-cross-entropy signature from Loss and Gradients Blog, now living in embedding space. When prediction matches target for every word, the gradient is zero and the vectors stop moving. That is what convergence looks like here.

Only the rows you touched get updated

One mechanical point, and it ties straight back to the lookup identity from §2. On a given training pair, the center vector is vc=Eec\mathbf{v}_c = E^\top \mathbf{e}_c ,the forward pass touches only row cc of the table EE. So when the gradient flows back, only row cc of EE receives a nonzero gradient; every other row gets exactly zero.

Why every other row gets zero gradient

Differentiate the loss with respect to an arbitrary entry EmjE_{mj}. Since vc=Eec\mathbf{v}_c = E^\top\mathbf{e}_c has components (vc)j=Ecj(v_c)_{j} = E_{cj}, the loss depends on EmjE_{mj} only when m=cm = c:

LEmj={L(vc)jm=c0mc\frac{\partial L}{\partial E_{mj}} = \begin{cases} \dfrac{\partial L}{\partial (v_c)_j} & m = c \\[4pt] 0 & m \neq c \end{cases}

Intuitively: a row the forward pass never read can't have influenced the loss, so nudging it changes nothing so its gradient is zero by definition.

This is why, out of a 50,000-row table, a single training step updates only the handful of rows whose words actually appeared. (It's also why deep-learning libraries have special "sparse gradient" paths for embedding layers because writing zeros to 49,990 untouched rows would be pure waste.)

One catch, and the fix

To compute p(oc)p(o \mid c), the probability of a single context word, the denominator

w=1Vexp(uwvc)\sum_{w=1}^{V} \exp(\mathbf{u}_w^\top \mathbf{v}_c)

sums over every word in the vocabulary. Think about what that means for one training pair like (ruled, queen): to nudge the vectors even slightly, the model computes a score for ruled against every one of the 100,000+ words in the language ("car," "photosynthesis," "the," all of them) just to normalize a single probability. And it does this for every pair, across billions of pairs. We are paying to score the entire dictionary just to learn one neighbour.

That cost is the denominator, and it also means every row of the context table UU gets touched on every step (all those uw\mathbf{u}_w appear in the sum), so unlike the center table EE, which only ever touched one row, UU never catches the sparse-gradient break. The full softmax is simply too expensive to run on a real vocabulary.

Here is the fix, and it comes from asking a slightly different question. The full softmax asks "which of all 100,000 words is the context?", a huge multiple-choice problem. Negative sampling asks the far cheaper question: "is this pair real, or did I make it up?" For the true pair (ruled, queen), we say "real" and pull vruled\mathbf{v}_{\text{ruled}} toward uqueen\mathbf{u}_{\text{queen}}. Then we invent a handful of fake pairs by grabbing random words, say (ruled, car), (ruled, and), (ruled, orange), label them "fake," and push vruled\mathbf{v}_{\text{ruled}} away from each. Five to fifteen random negatives, instead of the whole dictionary.

It is the same force field (pull the real neighbour close, shove a few random words away), just approximated with a tiny sample instead of the exact average over all VV words. And it turns an every-step cost of "score 100,000 words" into "score maybe 15," which is what makes word2vec fast enough to train on billions of tokens.

So, zooming out: embeddings learn by turning "predict your neighbours" into a loss, and letting gradient descent run a force field that drags co-occurring words together. What is genuinely surprising is what that space then contains: not just clusters of similar words, but linear structure rich enough that you can do arithmetic on meaning. That is the next section.

Arithmetic on Meaning

Once the embeddings are trained, you can do algebra on them:

EkingEman+EwomanEqueenE_{\text{king}} - E_{\text{man}} + E_{\text{woman}} \approx E_{\text{queen}}

Take the vector for king, subtract man, add woman, and you land almost exactly on queen. The same trick works for verb tense (walking - walked ++ swam \approx swimming), for capitals (Paris - France ++ Japan \approx Tokyo), for plurals, for comparatives. Somehow a relationship between words ("male to female," "country to capital") became a direction in the vector space, the same fixed step no matter which pair you apply it to.

Nobody programmed this in. We only ever told the model to predict neighbours.

The intuition: words are sums of their themes

Think about the contexts each word keeps. King shows up around royalty words (throne, reign, crown) and around male words (he, his, man). Queen shows up around the same royalty words but female ones instead. Man and woman show up around generic male and female contexts, without the royalty flavour.

Recall that a word's vector is shaped by the sum of the pulls from all its contexts. If a word lives in two kinds of context, its vector ends up as roughly the sum of two pieces. Write those pieces as components:

Ekingr+m,Equeenr+f,Emanm,EwomanfE_{\text{king}} \approx \mathbf{r} + \mathbf{m}, \qquad E_{\text{queen}} \approx \mathbf{r} + \mathbf{f}, \qquad E_{\text{man}} \approx \mathbf{m}, \qquad E_{\text{woman}} \approx \mathbf{f}

where r\mathbf{r} is the "royalty" component and m,f\mathbf{m}, \mathbf{f} are the "male" and "female" components. Now just do the arithmetic:

EkingEman+Ewoman(r+m)m+f=r+fEqueenE_{\text{king}} - E_{\text{man}} + E_{\text{woman}} \approx (\mathbf{r} + \mathbf{m}) - \mathbf{m} + \mathbf{f} = \mathbf{r} + \mathbf{f} \approx E_{\text{queen}}

The male component cancels, the female component takes its place, and the royalty component rides along untouched. The analogy works because the vectors are approximately additive in their underlying themes, and subtracting man while adding woman is exactly the operation "swap the gender theme, keep everything else."

The one thing this needs is that the gender contrast is consistent: the way male and female contexts differ has to be about the same for royals as for commoners, so that the "male to female" step is one fixed vector fm\mathbf{f} - \mathbf{m} regardless of which word you start from. In a large enough corpus, it statistically is, which is why the same offset resolves king\toqueen and actor\toactress alike.

This is the deep reason embeddings are more than a lookup table. A relationship that is consistent across many word pairs gets absorbed into a single direction, and once it is a direction, you can add and subtract it.

The rigorous version: word2vec is secretly matrix factorization

Levy and Goldberg (2014) showed what skip-gram with negative sampling is actually optimizing. At the optimum, the dot product of a context and center vector approximates a specific, classical quantity:

uovc    PMI(o,c)logk\mathbf{u}_o^\top \mathbf{v}_c \;\approx\; \text{PMI}(o, c) - \log k

where kk is the number of negative samples and PMI\text{PMI} is the pointwise mutual information between two words:

PMI(o,c)=logp(o,c)p(o)p(c)\text{PMI}(o, c) = \log \frac{p(o, c)}{p(o)\,p(c)}

PMI measures how much more often two words co-occur than you would expect if they were independent. It is large and positive for pairs like (doctor, nurse) that appear together far more than chance, near zero for unrelated pairs, and negative for pairs that avoid each other.

What this equation is saying

Stack the target dot products uovc\mathbf{u}_o^\top \mathbf{v}_c into a matrix MM with one entry per word pair. The result above says MPMIlogkM \approx \text{PMI} - \log k, a fixed matrix built purely from corpus co-occurrence counts.

But M=UVM = U V^\top is a product of two skinny matrices, each of width dd. A product of width-dd matrices has rank at most dd. So training skip-gram is finding a low-rank approximation of the (shifted) PMI matrix, which is exactly what techniques like SVD and PCA do. Word2vec is matrix factorization wearing a neural-network costume.

This is why the linear structure exists. The embeddings are a low-rank factorization of a matrix of co-occurrence statistics, so linear relationships among those statistics turn into linear (geometric) relationships among the vectors. The parallelogram king : queen :: man : woman holds precisely when the corresponding PMI relationships hold in the corpus, which in practice they approximately do. That is also why the analogy is an \approx and not an ==, and why it sometimes fails: it is only ever as clean as the underlying statistics.

Two honest caveats

Anyone who actually runs these analogies discovers two things the demos quietly skip.

First, you have to exclude the input words from the answer. When you compute EkingEman+EwomanE_{\text{king}} - E_{\text{man}} + E_{\text{woman}} and ask for the nearest vector, the true nearest neighbour is very often just king itself, because the offset you added is small compared to the distance between words. The standard evaluation removes the three query words from the candidate pool before picking the nearest. So the famous result is real, but it comes with an asterisk.

Second, the geometry absorbs whatever is in the corpus, including its biases. The same mechanism that learns king - man ++ woman \approx queen also learns doctor - man ++ woman \approx nurse, and programmer \to homemaker, straight from the statistical regularities of human-written text. The vectors are a faithful mirror of the corpus, which means they mirror its stereotypes too. This is not a bug in the math; it is the math working exactly as designed on biased data, and it is why de-biasing word embeddings became its own line of research.

Where this goes

Word2vec gives every word a single fixed vector. That is its ceiling: "bank" gets one point in space even though it means a riverbank in one sentence and a financial institution in another, and the single vector is stuck averaging the two. The next generation of models keeps the exact same starting object, an embedding table with the exact same lookup and sparse gradients from §2 and §3, but then transforms each vector based on the words around it, so that "bank" can land in different places depending on its sentence. That is the bridge from static embeddings to transformers, and it is the final section.