3. Word Embeddings

Meaning from context

Symbols and vectors

A language model built on counts treats every word as an atomic symbol. Written as vectors, symbols are one-hot: word \(i\) of a vocabulary of size \(V\) is the \(i\)-th standard basis vector of \(\mathbb R^V\), and any two different words are orthogonal and equally far apart. Nothing in this representation says that king is closer to queen than to carrot, which is why a trigram model learns nothing about dog from sentences about cat (chapter 2). A word embedding replaces the one-hot vector by a dense vector of a few hundred dimensions, \(e_w\in\mathbb R^d\), arranged so that words used alike lie close together. Every neural language model begins with such a table of vectors, and the methods of this chapter explain what those vectors can capture and why.

The guiding principle is the distributional hypothesis: words that occur in similar contexts tend to have similar meanings (Harris, 1954), or in Firth’s phrase, “you shall know a word by the company it keeps.” A reader who has never met the word ongchoi but sees that it is sautéed with garlic, served over rice, and has leaves like spinach can guess that it is a leafy vegetable. The methods below turn this idea into vectors in two ways, by counting contexts and factorizing the counts, and by training a model to predict contexts; the second turns out to be a version of the first.

Co-occurrence matrices

The most direct representation of a word’s contexts is a row of counts. In a term–document matrix, entry \((w,j)\) counts the occurrences of word \(w\) in document \(j\), so two words are similar if they occur in the same documents; this is the representation behind classical information retrieval (chapter 13). In a word–context matrix, entry \((w,c)\) counts how often context word \(c\) appears within a window of \(\pm L\) tokens around an occurrence of \(w\). The window size changes what similarity means. Small windows of one to three words capture functional similarity: words that can replace each other in a sentence, such as weep and mourn, or two first names. Large windows, or whole documents, capture topical relatedness: doctor and hospital, sword and battle.

Raw counts make poor vectors. They are dominated by frequent context words such as the and of, which co-occur with everything and say little about any particular word, and their scale grows with the frequency of \(w\).

Pointwise mutual information

Pointwise mutual information (Church and Hanks, 1990) measures how much more often \(w\) and \(c\) co-occur than they would if they were independent,

\[\operatorname{PMI}(w,c)=\log\frac{P(w,c)}{P(w)\,P(c)}=\log\frac{\#(w,c)\,|D|}{\#(w)\,\#(c)},\]

where \(\#(w,c)\) counts the co-occurrences, \(\#(w)=\sum_c\#(w,c)\) and \(\#(c)=\sum_w\#(w,c)\) are the marginal counts, and \(|D|\) is the total number of pairs. It is the log of the ratio in the mutual information of Foundations chapter 5, evaluated at a single pair of values rather than averaged. Positive values mark associations, zero independence, and negative values pairs that avoid each other; pairs never observed have \(\operatorname{PMI}=-\infty\). Negative values are unreliable, because showing that two words co-occur less than chance requires far more data than showing that they co-occur more, so the standard representation is positive PMI,

\[\operatorname{PPMI}(w,c)=\max\bigl(\operatorname{PMI}(w,c),0\bigr).\]

PMI overrates rare context words: a context seen once, next to \(w\), gets a large PMI from a single observation. A simple correction raises the context counts to a power \(\alpha<1\) before normalizing, \(P_\alpha(c)\propto\#(c)^\alpha\) with \(\alpha=0.75\), which raises the probability of rare contexts and lowers their PMI. The value comes from word2vec, and it helps count-based vectors as much as predictive ones (Levy, Goldberg, and Dagan, 2015).

Low-rank factorization

A PPMI matrix has as many columns as the vocabulary and is mostly zeros. Its truncated singular value decomposition \(M\approx U_d\Sigma_dV_d^\top\) keeps the \(d\) largest singular values, which gives the best rank-\(d\) approximation in Frobenius norm (Foundations chapter 2), and the rows of \(U_d\Sigma_d^{p}\) serve as \(d\)-dimensional word vectors, commonly with \(p=1/2\) or \(p=0\) rather than the “correct” \(p=1\). Applied to term–document matrices this is latent semantic analysis (Deerwester et al., 1990). The low rank does more than save memory: it forces words with similar but not identical contexts to share dimensions, so two words can become similar without ever sharing a context word, through words that share contexts with both.

The code counts co-occurrences within two words on each side in Tiny Shakespeare, for the 3,322 words that occur at least five times, converts them to PPMI with smoothed context probabilities, and keeps 100 dimensions of the SVD. It lists the six nearest neighbors of several words by cosine similarity.

import re
from collections import Counter

import numpy as np

text = open("Sources/Data/tinyshakespeare.txt", encoding="utf-8").read().lower()
toks = re.findall(r"[a-z]+(?:'[a-z]+)?", text)
counts = Counter(toks)
vocab = [w for w, c in counts.most_common() if c >= 5]
idx = {w: i for i, w in enumerate(vocab)}
V, window = len(vocab), 2
ids = np.array([idx.get(w, -1) for w in toks])

C = np.zeros((V, V))                                    # symmetric window co-occurrence counts
for off in range(1, window + 1):
    a, b = ids[:-off], ids[off:]
    ok = (a >= 0) & (b >= 0)
    np.add.at(C, (a[ok], b[ok]), 1)
    np.add.at(C, (b[ok], a[ok]), 1)

pw = C.sum(1) / C.sum()
pc = C.sum(0) ** 0.75 / (C.sum(0) ** 0.75).sum()       # smoothed context distribution
with np.errstate(divide="ignore"):
    pmi = np.log(C / C.sum()) - np.log(pw)[:, None] - np.log(pc)[None, :]
ppmi = np.maximum(pmi, 0)                               # positive PMI; -inf entries become 0
U, S, _ = np.linalg.svd(ppmi, full_matrices=False)
E = U[:, :100] * np.sqrt(S[:100])                       # 100-dimensional word vectors
E /= np.linalg.norm(E, axis=1, keepdims=True)
print(f"{len(toks)} tokens, {V} words with count >= 5, {int((C > 0).sum())} nonzero co-occurrences")
for w in ["king", "thou", "sword", "night", "romeo", "love", "weep"]:
    sims = E @ E[idx[w]]
    print(f"{w:6s}", " ".join(vocab[j] for j in np.argsort(-sims)[1:7]))
# 203839 tokens, 3322 words with count >= 5, 264272 nonzero co-occurrences
# king   richard henry edward vi york iii
# thou   art wilt dost hast not didst
# sword  hands date company gage promise ring
# night  morrow day last farewell romeo good
# romeo  juliet mercutio tybalt laurence nurse friar
# love   o hate soul juliet fear her
# weep   wherefore breathe choose oppression gaoler prevail

Even from 200,000 words the neighbors are sensible. Romeo is surrounded by the other characters of his play, thou by the verb forms that agree with it (art, wilt, dost), and king by the names that follow the word in King Richard and King Henry. The neighbors of sword and weep are poorer: both words are rarer, and a million characters of text do not pin down their contexts.

Learning vectors by prediction

The skip-gram model

Word2vec (Mikolov et al., 2013a; 2013b) learns vectors by prediction rather than counting. The skip-gram model gives every word two vectors, a word vector \(w\) used when it is the center and a context vector \(c\) used when it is in the window, and models the probability of seeing context \(c\) near word \(w\) as a softmax over the vocabulary,

\[P(c\mid w)=\frac{\exp(c^\top w)}{\sum_{c'}\exp(c'^\top w)}.\]

Training maximizes the log-likelihood of all (word, context) pairs within the window. The model is a log-bilinear language model with no hidden layer, and its only parameters are the two embedding tables. The companion continuous bag-of-words (CBOW) model predicts the center word from the average of its context vectors. The obstacle is the normalizer: each gradient step for one pair touches all \(V\) context vectors. The original papers avoided it with a hierarchical softmax over a binary tree of words, which costs \(O(\log V)\) per pair, and with the simpler method that became standard.

Negative sampling

Skip-gram with negative sampling (SGNS) replaces the softmax by binary classification. For each observed pair \((w,c)\) it draws \(k\) negative contexts \(c_1,\dots,c_k\) from a noise distribution \(P_n\) and maximizes

\[\log\sigma(c^\top w)+\sum_{i=1}^k\log\sigma(-c_i^\top w),\]

so the model learns to tell real pairs from random ones with a logistic regression whose features are the two vectors (ML chapter 5). The objective is a simplified form of noise-contrastive estimation (Gutmann and Hyvärinen, 2010) and a relative of the contrastive losses of DL chapter 10. Each step touches only \(k+1\) context vectors, with \(k\) between 2 and 20. The noise distribution is the unigram distribution raised to the power 0.75, the same smoothing as in PPMI. Word2vec also subsamples frequent words, discarding each occurrence of word \(w\) with probability \(1-\sqrt{t/f(w)}\) for its relative frequency \(f(w)\) and a threshold \(t\) around \(10^{-5}\), which speeds training and effectively widens the window around rare words.

What skip-gram computes

Although SGNS is trained by stochastic gradient descent on sampled pairs, its objective can be written as a sum over all word–context pairs with weights given by their counts, and each pair’s term depends only on the single number \(x=c^\top w\). If the dimension is large enough for every \(x\) to be set independently, the optimum is

\[c^\top w=\log\frac{\#(w,c)}{\#(w)\,P_n(c)}-\log k=\operatorname{PMI}(w,c)-\log k\]

(Levy and Goldberg, 2014; Appendix A), with PMI computed against the noise distribution. SGNS therefore factorizes the matrix of PMI values shifted by \(\log k\), like the SVD of the previous section but with a different loss: it weights each cell by how often the pair and its negatives were sampled, so frequent pairs are fit tightly and rare ones loosely, and it never sees the \(-\infty\) entries of unobserved pairs directly. The code trains SGNS in PyTorch on the same pairs as the PPMI code, with 100 dimensions and \(k=5\), then checks the prediction on the pairs seen at least 30 times.

import re
from collections import Counter

import numpy as np
import torch

torch.manual_seed(0)
rng = np.random.default_rng(0)
text = open("Sources/Data/tinyshakespeare.txt", encoding="utf-8").read().lower()
toks = re.findall(r"[a-z]+(?:'[a-z]+)?", text)
counts = Counter(toks)
vocab = [w for w, c in counts.most_common() if c >= 5]
idx = {w: i for i, w in enumerate(vocab)}
ids = np.array([idx[w] for w in toks if w in idx])
V, d, k, window = len(vocab), 100, 5, 2

# All (word, context) pairs within the window, and the noise distribution: unigram^0.75
pairs = np.concatenate([np.stack([ids[:-o], ids[o:]], 1) for o in range(1, window + 1)])
pairs = np.concatenate([pairs, pairs[:, ::-1]])
freq = np.bincount(ids, minlength=V).astype(float)
noise = torch.tensor(freq ** 0.75 / (freq ** 0.75).sum())

W = torch.nn.Embedding(V, d); Cv = torch.nn.Embedding(V, d)       # word and context vectors
torch.nn.init.uniform_(W.weight, -0.5 / d, 0.5 / d); torch.nn.init.zeros_(Cv.weight)
opt = torch.optim.Adam(list(W.parameters()) + list(Cv.parameters()), lr=3e-3)
P = torch.tensor(pairs)
for epoch in range(8):
    perm = torch.randperm(len(P))
    total = 0.0
    for s in range(0, len(P), 4096):
        w, c = P[perm[s:s + 4096], 0], P[perm[s:s + 4096], 1]
        neg = torch.multinomial(noise, len(w) * k, replacement=True).view(-1, k)
        vw = W(w)
        pos = torch.nn.functional.logsigmoid((vw * Cv(c)).sum(1))
        negs = torch.nn.functional.logsigmoid(-(Cv(neg) @ vw.unsqueeze(2)).squeeze(2)).sum(1)
        loss = -(pos + negs).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        total += loss.item() * len(w)
    print(f"epoch {epoch}: loss {total / len(P):.3f}")

E = torch.nn.functional.normalize(W.weight.detach(), dim=1)
for w in ["king", "thou", "sword", "night", "romeo", "love", "weep"]:
    sims = E @ E[idx[w]]
    print(f"{w:6s}", " ".join(vocab[j] for j in torch.argsort(-sims)[1:7].tolist()))

# At the optimum w.c = PMI(w, c) - log k, with P(c) the noise distribution.
pc = Counter(map(tuple, pairs.tolist()))
nw = np.bincount(pairs[:, 0], minlength=V)
keys = np.array([key for key, n in pc.items() if n >= 30])
n = np.array([pc[tuple(key)] for key in keys])
spmi = np.log(n / nw[keys[:, 0]]) - np.log(noise.numpy()[keys[:, 1]]) - np.log(k)
dots = (W.weight.detach()[keys[:, 0]] * Cv.weight.detach()[keys[:, 1]]).sum(1).numpy()
print(f"{len(keys)} pairs seen >= 30 times: correlation {np.corrcoef(spmi, dots)[0, 1]:.3f}, "
      f"mean dot {dots.mean():.2f} vs mean shifted PMI {spmi.mean():.2f}")
# epoch 0: loss 2.863
# epoch 1: loss 2.579
# epoch 2: loss 2.516
# epoch 3: loss 2.459
# epoch 4: loss 2.416
# epoch 5: loss 2.378
# epoch 6: loss 2.344
# epoch 7: loss 2.312
# king   crown convey'd foe marshal lewis richard
# thou   thyself liest couldst wilt wert dost
# sword  revenges conscience life company shoulders attempt
# night  day night's time goose morning bridal
# romeo  juliet mercutio banished laurence benvolio torch
# love   burden slander wits teach boldness juliet
# weep   conjure kneel mourn complain tarry intend
# 2476 pairs seen >= 30 times: correlation 0.980, mean dot -0.61 vs mean shifted PMI -0.63

For the 2,476 frequent pairs the learned dot products and the shifted PMI values correlate at 0.980, and their means differ by 0.02. The neighbors differ from those of PPMI and SVD in the direction the window suggests: weep now sits among verbs of the same kind (mourn, complain, kneel), thou among the verb forms it takes, and night next to day, time, and morning.

Skip-gram with negative sampling on Tiny Shakespeare (100 dimensions, \(k=5\), window of 2). Left: for the 2,476 word–context pairs seen at least 30 times, the learned dot product against \(\operatorname{PMI}(w,c)-\log k\) computed from the training pairs; the dashed line is the predicted optimum. Right: the correlation falls with pair frequency, from 0.98 for pairs seen 30 or more times to 0.88 for 5–9 times and 0.21 for pairs seen once, whose PMI estimates are mostly noise that 100 dimensions cannot and should not reproduce.

The equivalence explains a long debate. When word2vec appeared, its vectors beat count-based vectors on several benchmarks, which seemed to show that prediction was better than counting. Levy, Goldberg, and Dagan (2015) showed that most of the difference came from choices that could be transferred to either method, such as context-distribution smoothing, subsampling, the weighting of singular values, and adding the word and context vectors, and that with matched choices neither approach dominated.

GloVe

GloVe (Pennington, Socher, and Manning, 2014) makes the factorization explicit. It starts from the observation that ratios of co-occurrence probabilities carry meaning: \(P(\textit{solid}\mid\textit{ice})/P(\textit{solid}\mid\textit{steam})\) is large, \(P(\textit{gas}\mid\textit{ice})/P(\textit{gas}\mid\textit{steam})\) is small, and \(P(\textit{water}\mid\cdot)\) is about the same for both, so the relation between ice and steam shows in the ratios. Asking differences of vectors to encode log ratios leads to a weighted least-squares fit of the log counts,

\[J=\sum_{i,j:\,X_{ij}>0}f(X_{ij})\bigl(w_i^\top\tilde w_j+b_i+\tilde b_j-\log X_{ij}\bigr)^2, \qquad f(x)=\min\bigl(1,(x/x_{\max})^{3/4}\bigr),\]

where \(X_{ij}\) counts co-occurrences, \(b_i\) and \(\tilde b_j\) are biases that absorb the word frequencies, and the weight \(f\) discounts rare pairs and caps frequent ones at \(x_{\max}=100\). Only the nonzero cells enter the sum, so training costs time proportional to the number of distinct co-occurring pairs. GloVe and SGNS perform similarly in practice; both are factorizations of a transformed co-occurrence matrix with a weighting that favors frequent pairs.

Subword information

Word-level embeddings inherit the open-vocabulary problem of chapter 1: a rare word gets a poorly trained vector and an unseen word gets none. FastText (Bojanowski et al., 2017) represents a word as the sum of vectors for its character n-grams of lengths 3 to 6, with boundary markers, so where contributes <wh, whe, her, ere, re>, and longer pieces, plus a vector for the whole word. Words that share morphemes share parameters, an unseen word still gets a vector from its pieces, and the gains are largest for morphologically rich languages. Subword tokenizers made the same idea standard in later models, where each piece has its own embedding and the network composes them.

Using word vectors

Similarity and relatedness

Word vectors are compared by cosine similarity, \(\cos(u,v)=u^\top v/(\|u\|\|v\|)\), which ignores their lengths; lengths grow with word frequency and say more about how often a word appears than about what it means.

Sixty-five words from six hand-chosen categories, embedded by skip-gram in 100 dimensions and mapped to the plane with t-SNE (perplexity 8, cosine distance). In the full 100 dimensions the nearest neighbor of a word belongs to its own category for 89% of the words with skip-gram and 94% with PPMI and SVD. Some placements reflect the corpus: Montague sits among the noble houses, and grief and sorrow sit next to heart and blood.

Intrinsic evaluations correlate the cosine similarities of word pairs with human judgments, using Spearman’s rank correlation, on datasets such as WordSim-353, which mixes similarity with relatedness, and SimLex-999, which rates similarity alone and penalizes related but dissimilar pairs such as cup and coffee (Hill, Reichart, and Korhonen, 2015). Such benchmarks are cheap but correlate imperfectly with usefulness in downstream tasks, which is the evaluation that matters.

Analogies

Word vectors became famous for analogies. Answering “man is to woman as king is to ?” by the vocabulary word nearest to \(e_{\textit{king}}-e_{\textit{man}}+e_{\textit{woman}}\), excluding the three question words, returns queen for vectors trained on large corpora, and similar offsets encode plurals, verb tenses, comparatives, and capitals of countries. The offsets work when the relation changes the contexts of each pair of words in the same way: if every context is as much more likely with woman than with man as it is with queen than with king, then the PMI vectors satisfy the analogy exactly (Appendix B).

The phenomenon is more fragile than its reputation. Excluding the question words matters, because the nearest vector to \(e_{\textit{king}}-e_{\textit{man}}+e_{\textit{woman}}\) is usually king itself, and simple baselines that ignore part of the question, such as returning the nearest neighbor of the third word, already solve a substantial share of some benchmark categories (Linzen, 2016). Analogies also need large corpora. On Tiny Shakespeare, skip-gram answers two of eight simple family and gender analogies (he:she::his:her and father:mother::son:daughter) and PPMI with SVD none, although both place the words of each category together.

Social biases

Embeddings learn whatever associations the corpus contains, including stereotypes. Vectors trained on news text complete “man is to computer programmer as woman is to ?” with homemaker (Bolukbasi et al., 2016), and a test modeled on the psychologists’ implicit association test recovers, from common web text, the same associations between names, genders, and races on one side and pleasant or career-related words on the other that people show (Caliskan, Bryson, and Narayanan, 2017). Projecting out a “gender direction” removes the most visible effects, but the bias remains recoverable from the geometry of the other dimensions (Gonen and Goldberg, 2019). The same problem, at larger scale and with more ways to surface, runs through the language models of later chapters.

Word vectors inside larger models

Before pretrained language models, word vectors were the main vehicle of transfer learning in NLP: vectors trained on billions of words initialized, or were frozen as the input layer of, models for tagging, parsing, and classification trained on small labeled sets. A neural language model contains the same object. Its input embedding table maps each token to a vector, its output layer scores each token by a dot product with a second table of vectors (chapter 4), and both are trained end to end with the rest of the network, like the word and context vectors of skip-gram.

From static to contextual representations

One vector per word type

A static embedding gives each word type one vector, but many words have several meanings: bank of a river and of money, spring the season and the coil. Its vector is a mixture of its senses weighted by their frequencies, and the senses can be recovered from it: Arora et al. (2018) showed that the vector of a polysemous word lies close to a weighted sum of vectors for its senses, which sparse coding can separate. A mixture still cannot tell which sense is meant in a given sentence, and downstream models had to work that out from the surrounding vectors.

Contextual embeddings

A contextual embedding is a vector for a token in its sentence, computed by a network that reads the whole sentence. ELMo (Peters et al., 2018) took the hidden states of a bidirectional LSTM language model as word representations and gave large gains across tasks; BERT and GPT did the same with transformers, and the representation of a word in a pretrained transformer depends on everything around it (chapter 5). The static embedding survives as the first layer of these networks, which the upper layers progressively contextualize.

Sentences and documents

Averaging the word vectors of a sentence gives a surprisingly strong sentence representation, especially when frequent words are downweighted and the common direction shared by all sentences is removed (Arora, Liang, and Ma, 2017). Current sentence and document embeddings come from transformers fine-tuned with contrastive objectives on pairs of related texts (chapter 5), and they power the dense retrieval of chapter 13.

Appendices