How can a computer learn that great and excellent mean nearly the same thing? One answer is to learn word vectors by training them to solve a prediction task: given a word, guess which words appear around it. Words that appear in similar contexts will have to make similar predictions, so they will end up with similar vectors.
This is the idea behind Word2Vec (Mikolov et al., 2013), one of the most influential ideas in NLP. In this reading, we'll set up its main model, skip-gram: we'll define the prediction task, turn word vectors into probabilities, and write down the objective function that training minimizes.
1What we want from word vectors
The goal is to map each word to a short, dense vector of real numbers, say $d=100$ to $300$ of them, so that words with similar meanings get similar vectors. Similarity is measured with cosine:
$$\text{similarity}(w_1,w_2)=\cos(\mathbf v_{w_1},\mathbf v_{w_2})=\frac{\mathbf v_{w_1}\cdot\mathbf v_{w_2}}{\|\mathbf v_{w_1}\|\,\|\mathbf v_{w_2}\|}$$Nobody writes these numbers by hand. The question for this reading is: how can a computer learn them from raw text?
2Word2Vec's prediction task
Word2Vec rests on the distributional hypothesis: a word is characterized by the company it keeps. Look at a few occurrences of banking:
…government debt problems turning into banking crises as happened in 2009…
…saying that Europe needs unified banking regulation to replace the hodgepodge…
…India has just given its banking system a shot in the arm…
The word in the middle is the center word, and the words around it are its context words. The relationship between them can be turned into a prediction problem in two directions:
- Given the context words, predict the center word. This is CBOW (continuous bag of words).
- Given the center word, predict the context words. This is skip-gram, the main model in this reading.
Concretely, skip-gram slides through the corpus one position at a time. At each position $t$, the word $w_t$ is the center word, and every word within a fixed window of $m$ positions on either side, $w_{t-m},\dots,w_{t-1},w_{t+1},\dots,w_{t+m}$, is a context word. Each (center, context) pair becomes one training example. Click any word below to make it the center, and change the window size:
Notice what this asks of the model. When into is the center, it should assign high probability to turning and banking. When banking is the center, it should assign high probability to into and crises. Every word must be able to say, for each word in the vocabulary, how likely that word is to show up nearby. So skip-gram is a classification problem: for a given center word, it outputs a probability distribution over all $|V|$ possible context words.
3From vectors to probabilities
Step 1: a score from a dot product
How should a pair of word vectors produce $P(w_\text{context}\mid w_\text{center})$? The simplest idea is to use the dot product as a score: if two vectors point in similar directions, the score is high, and the context word should be more likely. Scores can be any real number, though, and probabilities must be non-negative and sum to 1. The softmax function fixes that:
$$\text{softmax}(\mathbf z)_i=\frac{e^{z_i}}{\sum_{j} e^{z_j}}$$Exponentiating makes every score positive, and dividing by the sum makes them add up to 1. Larger scores still get larger probabilities, and the largest score gets amplified the most (which is where the name “soft max” comes from). Applying softmax to the scores of every word in the vocabulary gives a first attempt:
$$P(w_\text{context}\mid w_\text{center})=\frac{\exp(\mathbf v_{w_\text{center}}\cdot\mathbf v_{w_\text{context}})}{\sum_{k\in V}\exp(\mathbf v_{w_\text{center}}\cdot\mathbf v_k)}$$Try it below. Each slider sets one dot-product score for the center word banking; the bars show the resulting probabilities.
Step 2: two vectors per word
The one-vector version has two problems. First, the score of a word with itself is $\mathbf v_w\cdot\mathbf v_w=\|\mathbf v_w\|^2$, which is always non-negative and usually large. The model would think banking is very likely to appear right next to banking, yet a word rarely appears in its own context. (Switch Figure 3 to “one vector” to see this.) Second, the dot product is symmetric, so $P(b\mid a)$ and $P(a\mid b)$ would share the same score, but the two relationships need not be equally strong. The scores should be asymmetric.
The fix is to give every word two vectors:
- $\mathbf u_w\in\mathbb R^d$: the vector of $w$ when it is the center word,
- $\mathbf v_w\in\mathbb R^d$: the vector of $w$ when it is a context word.
The final skip-gram model is then
$$P(w_{t+j}\mid w_t;\theta)=\frac{\exp(\mathbf u_{w_t}\cdot\mathbf v_{w_{t+j}})}{\sum_{k\in V}\exp(\mathbf u_{w_t}\cdot\mathbf v_k)}$$The numerator scores how well the context word $w_{t+j}$ fits the center word $w_t$; the denominator normalizes over the entire vocabulary to give a probability distribution. Here $\theta$ stands for all the parameters, namely every $\mathbf u_w$ and every $\mathbf v_w$: that's $2\,|V|\,d$ numbers in total, for example $2\times100{,}000\times300=60$ million. After training, people usually keep $\mathbf u_w$ as the word vector, or use the average $(\mathbf u_w+\mathbf v_w)/2$.
4The training objective
We now have a model that assigns a probability to every (center, context) pair. Training should make the pairs we actually observed as likely as possible. For a corpus of $T$ words, the likelihood multiplies the probability of every observed context word, at every position:
$$\mathcal L(\theta)=\prod_{t=1}^{T}\;\prod_{\substack{-m\le j\le m\\ j\neq 0}} P(w_{t+j}\mid w_t;\theta)$$The outer product runs over every position in the corpus; the inner product runs over every context word in the window around that position. For the sentence in Figure 2 with $m=2$, the into and banking positions alone contribute
$$P(\text{problems}\mid\text{into})\,P(\text{turning}\mid\text{into})\,P(\text{banking}\mid\text{into})\,P(\text{crises}\mid\text{into})\;\times\;P(\text{turning}\mid\text{banking})\,P(\text{into}\mid\text{banking})\,P(\text{crises}\mid\text{banking})\,P(\text{as}\mid\text{banking})\cdots$$Products of millions of probabilities are numerically unusable (they underflow to 0), and maximizing is less conventional than minimizing. So we take the log, which turns products into sums, flip the sign, and average over positions. The result is the objective function (also called the loss or cost), the average negative log-likelihood:
$$J(\theta)=-\frac{1}{T}\log\mathcal L(\theta)=-\frac{1}{T}\sum_{t=1}^{T}\sum_{\substack{-m\le j\le m\\ j\neq0}}\log P(w_{t+j}\mid w_t;\theta)$$Maximizing the likelihood and minimizing $J(\theta)$ are the same thing, because $\log$ is increasing. In this reading, $\log$ is the natural log.
5Putting it all together
The complete skip-gram model has three parts:
| Part | What it is |
|---|---|
| Training data | every (center, context) pair within a window of $\pm m$ words |
| Parameters $\theta$ | two vectors per word: $\mathbf u_w$ (center) and $\mathbf v_w$ (context), $2\,|V|\,d$ numbers in total |
| Model | $P(w_{t+j}\mid w_t;\theta)=\text{softmax}$ over $\mathbf u_{w_t}\cdot\mathbf v_k$ for all $k\in V$ |
| Objective | minimize the average negative log-likelihood $J(\theta)$ |
Two properties of this objective are worth keeping in mind. First, $J(\theta)$ is a smooth function of every $\mathbf u_w$ and $\mathbf v_w$, so it can be minimized with standard gradient-based optimization. Second, $J(\theta)$ can never reach 0. A center word like banking appears with many different context words, and the model has to spread its probability across all of them. The best the model can do is to match how often each context actually occurs.
Nothing in the objective says that “similar words should have similar vectors.” That property emerges on its own: if bread and cake appear with the same context words (bake, sweet, eat, …), then the cheapest way to predict those contexts well is to give $\mathbf u_{\text{bread}}$ and $\mathbf u_{\text{cake}}$ similar directions.
6Check your understanding
Exercise: softmax probabilities by hand
The vocabulary has three words $\{a,b,c\}$ and $d=2$. The center word has $\mathbf u_o=[1,0]$, and the context vectors are $\mathbf v_a=[2,0]$, $\mathbf v_b=[0,1]$, $\mathbf v_c=[-1,0]$. Compute (1) $P(a\mid o)$, $P(b\mid o)$, $P(c\mid o)$; (2) the loss $-\log P(\cdot\mid o)$ if the observed context word is $a$; (3) the same loss if the observed context word is $b$ instead.
Exponentials: $e^2\approx7.389$, $e^0=1$, $e^{-1}\approx0.368$, with sum $\approx8.757$.
(1) $P(a\mid o)\approx0.844$, $P(b\mid o)\approx0.114$, $P(c\mid o)\approx0.042$. They sum to 1.
(2) $-\ln 0.844\approx0.170$.
(3) $-\ln 0.114\approx2.17$.
The loss is small when the model already gives the observed word high probability, and large when it doesn't. Training adjusts $\mathbf u$ and $\mathbf v$ to lower the total loss over all observed pairs.
Quick quiz
7Further reading
- Jurafsky & Martin, Speech and Language Processing (3rd ed. draft), Ch. 6, “Vector Semantics and Embeddings,” Section 6.8 (Word2vec).
- Mikolov, Chen, Corrado & Dean (2013). Efficient Estimation of Word Representations in Vector Space. ICLR Workshop.
- Rong (2014). word2vec Parameter Learning Explained. arXiv:1411.2738.