Word2Vec: skip-gram and CBOW

The previous lesson, on dense representations, built the table and left it empty: the embedding matrix $\mathbf{E}$ has a shape, a size and a geometry you can now read, and not one number in it. The distributional hypothesis asserts where meaning lives and does not say how to compute it. Word2Vec was the first to close that gap at real scale, in 2013, and it comes from an unexpected direction: it computes no one's coordinates, it lets them fall out as the residue of another task.

The task is a guessing game. Take the dog barks in the yard of the school, cover everything but dog and ask what was around it; a model that answers barks before school has learned something, and to answer it needs numbers to score candidates with. Those numbers are in the matrix $\mathbf{E}$, which enters a model with random values and gets corrected until it gets them right (that is what training is), after which the model is thrown away and $\mathbf{E}$ is kept. Word2Vec is that idea turned into a technique, and it is three pieces: a prediction task over windows of text, two ways of solving it (skip-gram and CBOW, continuous bag of words) and a simplification that makes the computation viable: negative sampling. What is left to decide is which parameters do the scoring, and what "getting it right" means once it is measured by a count that, written out with care, comes out unaffordable.

Slide a window of $m$ tokens on each side along the corpus: the middle one is the centre, the ones at the sides are its context, and one pair comes out for each neighbour. With $m = 2$, the centre yard of the sentence above produces four: with in, the, of and the again. The raw corpus already carries the pairs inside it: the distributional hypothesis's promise, collected.

Predicting the context: the skip-gram objective

Fix the window at $m$ and walk the corpus. The skip-gram takes each position $t$, with its central token $w_t$, and asks to predict each of its neighbours separately. Getting it right is written as a loss: a number, computed from the parameters and the whole corpus, that measures how badly the model does; correcting them until it gets it right is making that number small. The skip-gram's is the mean negative log-likelihood of the pairs:

where $t$ runs over the $T$ positions of the corpus and the minus sign turns the criterion around: the probability of the real pair has to go up, so $\mathcal{L}$ has to go down. The essential thing is missing: what $P(c \mid w)$ is.

Predicting is choosing among candidates, and choosing requires scoring them. We need a number that says how well the centre $w$ and a candidate context $c$ fit together. Word2Vec takes the simplest one between two vectors, their dot product: it grows when they point in similar directions and when they are long. We call it the pair's affinity.

For now we have only the table $\mathbf{E}$ with one vector per entry; the affinity of a pair comes from reading two of its rows. That forces every entry to use the same vector as centre and as candidate, and the flaw shows in the pair of an entry with itself: $\mathbf{e}_w^{\top}\mathbf{e}_w = \lVert\mathbf{e}_w\rVert^2$, positive and larger the longer the vector is. The cosine of the previous lesson cancelled it by dividing by the norms; in the affinity there is no division. And the corpus asks for the opposite: dog almost never falls inside its own window, and it has to be possible to score that pair low. But lowering $\lVert\mathbf{e}_w\rVert^2$ forces $\mathbf{e}_w$ to shrink, and that lowers every affinity in which $w$ is the centre at the same time.

Hence the design decision that defines Word2Vec: each entry receives two vectors of $d_{\text{model}}$ coordinates, the $\mathbf{e}_w$ for when it acts as centre and a new one, $\mathbf{u}_c$, for when it acts as context. They live in two tables of the same shape, $\mathbf{E}, \mathbf{U} \in \mathbb{R}^{\lvert V \rvert \times d_{\text{model}}}$, with row $i$ of $\mathbf{U}$ equal to $\mathbf{u}_{w_i}^{\top}$ exactly as that of $\mathbf{E}$ is $\mathbf{e}_{w_i}^{\top}$. The affinity becomes $\mathbf{u}_c^{\top}\mathbf{e}_w$, and that of an entry with itself, $\mathbf{u}_w^{\top}\mathbf{e}_w$, can go down without $\mathbf{e}_w$ shrinking.

Reading affinity as a probability asks for it to be made positive and for all of them to sum to 1; the softmax does both:

The exponential function preserves order, so the candidate with the most affinity takes the highest probability. And keep this about the denominator: it is there so that the $\lvert V \rvert$ affinities are a probability distribution over $V$, and for nothing else.

The parameters are $\theta = (\mathbf{E}, \mathbf{U})$, and $\mathcal{L}$ depends on nothing else. To minimise it we use gradient descent, which we will see in detail in Block 2, on the multilayer perceptron.

The gradient $\nabla_{\theta}\mathcal{L}$ points in the direction of the loss's steepest increase. So, at each step $t$, we update the parameters in the opposite direction,

where the subscript $t$ counts descent steps, not corpus positions, and $\eta$ is the learning rate, the size of the step. Choosing it is that block's business: too large and the descent diverges.

CBOW: the same idea with the roles swapped

CBOW inverts the question: it guesses the centre from the context. The $2m$ neighbours enter averaged into a single vector,

and the probability of the centre is the same softmax with $\bar{\mathbf{e}}_t$ in place of $\mathbf{e}_w$. The average loses the order: the bag of words of the lesson on TF-IDFBlock 1 · Lesson 5Bag of words and TF-IDFSumming the one-hot vectors of a document gives a count vector where the function words rule. Weighting each coordinate by how rare the entry is in the rest of the corpus fixes that, and the logarithm that does it is no trick: it comes from measuring information. in miniature, and hence the name.

Both are used: CBOW is faster (one prediction per position instead of $2m$), and skip-gram makes more of the rare entries, which lose out when they are averaged in. What they share is what matters now: a softmax over $V$.

The sum that makes the objective unaffordable

Count the operations of a single probability. The numerator is a dot product, $d_{\text{model}}$ multiplications; the denominator, one per vocabulary entry:

Fifteen million for a single pair. With $2m = 4$ pairs per token and a corpus of $T = 10^{10}$, that gives $4 \times 10^{10}$ pairs, and a whole pass costs $6 \times 10^{17}$ operations, and training does several. No machine was going to pay that for an embedding, not in 2013 and not now.

Look at where the cost is, because that is where the way out comes from. Making barks score high next to dog costs 300 multiplications; the other 14,999,700 go on a sum whose only job is to make the probabilities sum to one. You pay for a distribution over 50,000 entries you are not going to use: when it ends you throw the model away and keep $\mathbf{E}$.

So the exact objective is abandoned, and it is worth saying so in full, because it is a move the course will make again, in Block 5, with attention: when the right quantity is unaffordable, you swap it for one you can afford with similar optima, and then check whether the swap broke anything. What you do not do is pretend the second is the first.

Negative sampling: from sharing probability out to telling pairs apart

The substitution changes the question. In place of asking, given dog, what probability each of the 50,000 entries has of going with it, it asks a different question: is this pair, drawn at random, from the corpus, or did I make it up?

False pairs are needed, and they are manufactured: to the real pair $(w, c)$ we add $n_{\text{neg}}$ pairs $(w, c_j)$ with the context drawn at random. The original paper fixed the number by trying: from $5$ to $20$ with small corpora, from $2$ to $5$ with large ones. Take $n_{\text{neg}} = 5$. Each pair's probability is computed by squashing its affinity into the interval $(0,1)$ with the logistic function

The loss of a real pair and its negatives is

and it reads left to right. The first term goes down when $\sigma(\mathbf{u}_c^{\top}\mathbf{e}_w)$ rises towards $1$: when the centre moves closer to the real context. In the second, $\sigma(-x)$ is the probability of "no", and each summand goes down when the affinity with an invented context turns negative. Move closer to one, away from the other five.

The negatives are not drawn uniformly at random, but from

with $f_w$ the frequency of the type $w$ in the corpus. The exponent $3/4$ compresses the distribution without flattening it, between two bad extremes: with $1$ the negatives would be the function words of the lesson on the bag of wordsBlock 1 · Lesson 5Bag of words and TF-IDFSumming the one-hot vectors of a document gives a count vector where the function words rule. Weighting each coordinate by how rare the entry is in the rest of the corpus fixes that, and the logarithm that does it is no trick: it comes from measuring information., with $0$ the rarest entries, and neither costs anything to get right. There is no derivation behind it; it is the value that worked.

Now the count, the reason for all of this. A step touches the centre, the real context and the $n_{\text{neg}}$ negatives, one dot product per vector:

against the 15 million of before. The factor is $\lvert V \rvert / (n_{\text{neg}} + 1) \approx 8\,333$, and the whole pass drops from $6 \times 10^{17}$ to $7.2 \times 10^{13}$: from impossible to an afternoon. Notice what has disappeared and what has not: the cost per pair stops depending on $\lvert V \rvert$, and the $4 \times 10^{10}$ pairs are still there.

The concession, out loud: this is not the softmax computed faster, it is a different loss, and raising $n_{\text{neg}}$ does not make it converge to $\mathcal{L}$. Whether its optima place the vectors where they serve us is checked by looking at the neighbours that come out. The paper brings a second shortcut with the same end, hierarchical softmax, which this lesson does not go through.

The step needs $\partial \mathcal{L}_{\text{par}} / \partial \mathbf{e}_w$, and it comes from two derivatives that say, in passing, where each term pushes.

Training a tiny skip-gram in NumPy

The cell trains for real, in miniature: 800 tokens of a toy corpus from templates, 20 entries, $d_{\text{model}} = 12$, $n_{\text{neg}} = 5$, four passes. The update is the gradient from the derivation above: g = s - labels, where s is the output of the sigmoid applied to each pair's affinity and labels is the vector with 1 for the real pair and 0 for the negatives. Subtracting it, scaled, from each vector is gradient descent. The learning rate starts at $0.05$ and decays to almost zero: the last steps refine instead of jumping. Run it and look at the neighbours of dog before and after.

Before training, the nearest neighbour to dog is kitchen at $0.5$: pure noise. Four passes later they are cat, mouse and horse, all at least $0.96$; meat's are the other three foods and yard's the other three places. No one has said they resemble each other: only window pairs have gone in, and a $1$ or a $0$ per pair.

Two things, unvarnished. The corpus is artificial and the animals are interchangeable by construction: hence cosines that graze $1$ and do not come out with real text. And girl lands among the animals, next to dog and cat: it appears in a single mould (the subject of sees, the same slot the animals fill), and it is the mould that got learned, not the word.

Set n_neg = 0 and run it again. The loss collapses to $0.019$ (with no negatives, saying yes to everything is perfect), and the neighbours turn to junk: cosines pinned near $0.99$ across the board, dog now as close to the as to another animal. The negatives were not only the saving: with nothing to push against, the loss goes down and the representation does not exist.

Test your intuition

Five questions: what is kept at the end, how many pairs a window gives, where the cost was, what negative sampling changes and what the exponent $3/4$ does.

The matrix $\mathbf{E}$ is no longer empty; the debt the previous lesson ended on is paid. One peculiarity remains. Word2Vec does not see the corpus, it sees one window at a time, and it only concludes that dog and cat resemble each other by stumbling into thousands of those windows. But that is already told before any training: a single pass noting how many times each entry fell inside another's window gives a $\lvert V \rvert \times \lvert V \rvert$ table of co-occurrences where the row of dog and that of cat already resemble each other. Word2Vec never gets around to writing it.

Making use of it is the lesson on GloVe and the limits of static embeddings: GloVe (global vectors) builds that table and looks for the vectors that explain it, another objective with results of the same family. And with the two techniques on the table, that lesson closes the block by charging what the word "static" hides: each entry has one vector, a single one, for ever.