Learning-Based Word Vectors: Training and Variants

The skip-gram model of Word2Vec learns word vectors by predicting the words that appear around each word. Writing down that model is only half the story. We still have to find vectors that actually make good predictions, and do it fast enough to handle billions of words of text.

This reading is about training. We'll derive the gradient of the skip-gram loss, run stochastic gradient descent, and see why negative sampling is needed to make training practical. Then we'll look at three related methods that change what is predicted or what is learned from: CBOW, GloVe, and FastText.

1Recap: the skip-gram model

Training data. Slide through the corpus. At each position $t$, the word $w_t$ is the center word, and every word within $m$ positions of it is a context word. Each (center, context) pair is one training example:

Model. Every word $w$ has two vectors: $\mathbf u_w\in\mathbb R^d$ when it is the center word, and $\mathbf v_w\in\mathbb R^d$ when it is a context word. The probability of a context word given a center word is a softmax over dot products:

Objective. The parameters $\theta=\{\mathbf u_k\},\{\mathbf v_k\}$ are chosen to minimize the average negative log-likelihood of all observed pairs in a corpus of $T$ words:

Throughout this reading, $\log$ is the natural log.

2Training with gradient descent

Stochastic gradient descent

We can't solve $\min_\theta J(\theta)$ in closed form, but we can improve $\theta$ step by step. The gradient $\nabla_\theta J(\theta)$ points in the direction in which $J$ increases fastest, so we move a little in the opposite direction:

The learning rate $\eta$ controls the step size. Too small and training crawls; too large and the updates overshoot. Computing the full gradient needs a pass over the entire corpus, which is far too slow for one step. Stochastic gradient descent (SGD) instead estimates the gradient from a single training pair (or a small batch of pairs) and updates right away. Each estimate is noisy, but there are many of them and they are cheap.

The gradient for one pair

Since $J(\theta)$ is a sum of terms $-\log P(w_{t+j}\mid w_t;\theta)$, its gradient is the sum of the gradients of the terms. So it is enough to work out the gradient for one pair. Following the slides, write $o$ for the center word and $c$ for the context word, and let

Differentiate with respect to the center vector $\mathbf u_o$. The first term gives $-\mathbf v_c$. For the second, use $\frac{d}{dx}\log x=\frac1x$ and $\frac{d}{dx}e^x=e^x$ (the chain rule):

This has a lovely reading. $\mathbf v_c$ is the context vector we observed, and $\sum_k P(k\mid o)\,\mathbf v_k$ is the context vector the model expected, the average of all context vectors weighted by the model's current probabilities. The gradient is “expected minus observed”, so the update $\mathbf u_o\leftarrow\mathbf u_o-\eta\,\partial y/\partial\mathbf u_o$ pulls $\mathbf u_o$ toward the observed context and away from what the model currently predicts. When the model already puts all its probability on $c$, the two cancel and nothing changes.

A similar calculation gives the gradient for each context vector:

For the true context word ($k=c$), the coefficient $P(c\mid o)-1$ is negative, so $\mathbf v_c$ moves toward $\mathbf u_o$. For every other word, the coefficient $P(k\mid o)$ is positive, so $\mathbf v_k$ is pushed away from $\mathbf u_o$, more strongly the more probability the model wrongly gave it.

Putting it together

1. Randomly initialize all $\mathbf u_i$ and $\mathbf v_i$ to small values.
2. Walk through the corpus and collect (center, context) pairs $(o,c)$ as in Figure 1.
3. For each pair, update $\mathbf u_o\leftarrow\mathbf u_o-\eta\,\dfrac{\partial y}{\partial\mathbf u_o}$ and $\mathbf v_k\leftarrow\mathbf v_k-\eta\,\dfrac{\partial y}{\partial\mathbf v_k}$ for all $k\in V$.
4. Repeat for several passes (epochs) over the data.

The figure below runs exactly this algorithm, live in your browser, on a tiny corpus of 120 sentences such as “i bake the apple sweet” and “we install the laptop run”. To make the vectors directly plottable, it uses $d=2$. The plot shows each word's center vector $\mathbf u_w$.

Watch the loss curve. The loss drops quickly at first and then levels off. It can't reach 0, because the task is inherently uncertain: after bake, several food words are equally plausible. Then reset and try a large learning rate such as 0.5: the updates overshoot, and the loss stays high and bounces around instead of settling.

3Negative sampling

There is a serious efficiency problem hidden in the gradient. The softmax denominator sums over the entire vocabulary, so every single training pair requires computing $|V|$ dot products and updating all $|V|$ context vectors $\mathbf v_k$. With $|V|=100{,}000$ and billions of training pairs, that's hopeless.

Negative sampling (Mikolov et al., 2013b) replaces the $|V|$-way classification with a much cheaper question: “is this pair a real (center, context) pair, or a fake one?” For each real pair $(o,c)$, we draw $K$ random “negative” words $n_1,\dots,n_K$ (typically $K=5$ to $20$) that probably did not appear near $o$. The loss becomes

The sigmoid $\sigma$ squashes a score into a probability between 0 and 1. The first term wants $\sigma(\mathbf u_o\cdot\mathbf v_c)$ close to 1: the real pair should get a high score. Since $\sigma(-x)=1-\sigma(x)$, the second term wants each $\sigma(\mathbf u_o\cdot\mathbf v_{n_j})$ close to 0: the fake pairs should get low scores. Now each update touches only $\mathbf u_o$ and $K+1$ context vectors, instead of all $|V|$ of them: with $|V|=100{,}000$ and $K=5$, that is over 10,000 times less work per pair.

Which words are sampled?

The negatives are drawn from the unigram distribution raised to the $3/4$ power, $P_n(w)\propto\text{count}(w)^{3/4}$. Compared with sampling in proportion to raw frequency, this samples very frequent words like the a bit less often, and rare words a bit more often.

Negative sampling is what made Word2Vec fast enough to train on billions of words on a single machine. The variant described here is called skip-gram with negative sampling (SGNS), and it is the version most people mean when they say “Word2Vec.”

4CBOW: predicting the center word

Word2Vec actually comes in two flavors, which flip the direction of the prediction:

Continuous bag of words (CBOW) averages the context vectors in the window into one vector, and uses it to predict the center word:

It's called a “bag of words” because averaging throws away the order of the context words. CBOW makes one prediction per position instead of $2m$, so it trains faster. Skip-gram makes a separate prediction for every context word, so rare words receive more training signal; it tends to give better vectors for infrequent words, which is why it's the more common choice.

5GloVe: back to global counts

Word2Vec learns from one window at a time and never looks at the corpus-wide counts directly. GloVe (“Global Vectors,” Pennington et al., 2014) takes a different approach: it first builds the global co-occurrence matrix $X$, where $X_{ij}$ counts how often word $j$ appears in the context of word $i$, and then learns vectors that fit those counts.

The key observation is that ratios of co-occurrence probabilities carry meaning. Compare the words ice and steam by checking how often various probe words $x$ appear near each:

GloVe wants this ratio to be captured by vector differences. Its model is log-bilinear, $\mathbf w_i\cdot\tilde{\mathbf w}_j\approx\log P(i\mid j)$, which implies

so the direction $\mathbf w_{\text{ice}}-\mathbf w_{\text{steam}}$ encodes the “solid vs. gas” contrast. In practice, GloVe minimizes a weighted least-squares loss over all nonzero cells of $X$:

The biases $b_i$ and $\tilde b_j$ absorb how frequent each word is on its own. The weighting function $f$ (with $x_{\max}=100$ in the paper) gives rare, noisy co-occurrences little weight, and stops very frequent pairs like (the, of) from dominating. Pairs that never co-occur ($X_{ij}=0$) are skipped entirely, so training cost depends on the number of nonzero cells rather than the corpus size. That makes GloVe fast and able to scale to very large corpora.

6FastText: vectors for pieces of words

All the models so far treat each word as an atomic symbol. That causes two problems. A word never seen in training gets no vector at all, and related forms such as teach, teacher, and teaching are learned completely independently, even though they share most of their letters. This hurts most for rare words and for languages with rich morphology.

FastText (Bojanowski et al., 2017) keeps the skip-gram setup but represents each word by its character n-grams. The word is wrapped in boundary markers < and >, every character n-gram with $3\le n\le 6$ gets its own vector, and the word itself (with its markers) is added as one more unit. Type a word to see its pieces:

The score of a (center, context) pair is then the sum over all of the center word's units $\mathcal G(w)$:

Words that share n-grams share parameters, so teaching benefits from everything learned about teach and teacher. And a word never seen in training, like a typo or a new coinage, still gets a sensible vector by summing the vectors of its n-grams.

7Putting it all together

Skip-gram learns from local windows and GloVe from global counts, but the two are more closely related than they look. Levy & Goldberg (2014) showed that skip-gram with negative sampling implicitly factorizes a word–context matrix whose cells are $\text{PMI}(w,c)-\log K$, where $\text{PMI}(w,c)=\log\frac{P(w,c)}{P(w)\,P(c)}$ measures how much more often $w$ and $c$ co-occur than chance would predict. So both kinds of methods learn from the same underlying co-occurrence statistics.

You rarely need to train word vectors yourself. Pretrained vectors are freely available for all three methods: Word2Vec, GloVe, and FastText (which covers 157 languages). With the gensim library, loading and querying them takes a few lines:

import gensim.downloader as api

wv = api.load("glove-wiki-gigaword-100")        # 400k words, d = 100
wv.most_similar("banking", topn=5)               # nearest neighbors by cosine
wv.similarity("great", "excellent")                 # cosine similarity

# or train skip-gram with negative sampling on your own corpus
from gensim.models import Word2Vec
model = Word2Vec(sentences, vector_size=100, window=5, sg=1, negative=5)

One limitation remains: every word gets a single vector, no matter which sense it's used in. bank in “river bank” and in “savings bank” share one vector. Fixing that requires vectors that depend on the surrounding sentence.

8Check your understanding

Exercise: one gradient step 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]$. The observed context word is $a$. Compute (1) $P(a\mid o)$, $P(b\mid o)$, $P(c\mid o)$; (2) the loss $y=-\log P(a\mid o)$; (3) the gradient $\partial y/\partial\mathbf u_o$.

Quick quiz

9Further reading

• Jurafsky & Martin, Speech and Language Processing (3rd ed. draft), Ch. 6, “Vector Semantics and Embeddings,” Section 6.8 (Word2vec).
• Rong (2014). word2vec Parameter Learning Explained. arXiv:1411.2738. (Detailed gradient derivations.)
• Mikolov, Sutskever, Chen, Corrado & Dean (2013). Distributed Representations of Words and Phrases and their Compositionality. NeurIPS. (Negative sampling.)
• Mikolov, Chen, Corrado & Dean (2013). Efficient Estimation of Word Representations in Vector Space. ICLR Workshop. (CBOW and skip-gram.)
• Pennington, Socher & Manning (2014). GloVe: Global Vectors for Word Representation. EMNLP.
• Bojanowski, Grave, Joulin & Mikolov (2017). Enriching Word Vectors with Subword Information. TACL.
• Levy & Goldberg (2014). Neural Word Embedding as Implicit Matrix Factorization. NeurIPS.