Word embeddings

The big idea

In a bag of words ( Bag of words and TF-IDF ), every word is its own column. "great" and "excellent" are as unrelated as "great" and "terrible": a model that learned "excellent" is positive knows nothing about "great". What we want is a representation where similar words get similar numbers. A word embeddingA dense vector of real numbers, typically 50 to 300 dimensions, learned for each word so that words used in similar contexts get similar vectors. gives every word a short, dense vector, a few hundred numbers instead of a 50,000-wide one-hot column, and places words with similar meanings close together.

Where does meaning come from, if nobody labels it? From company. If you read "she poured the ___ into her cup" and "the ___ was too hot to drink", you can guess the blank is coffee, tea or soup without a dictionary. Words that appear in the same contexts tend to mean similar things. That is the distributional hypothesisThe idea that words occurring in similar contexts tend to have similar meanings, which lets meaning be learned from co-occurrence statistics., and every method in this lesson turns it into numbers.

How it works

Step 1: From context to vectors

Slide a window over a large corpus and record which words appear near which. "coffee" and "tea" share neighbours like "cup", "hot" and "drink"; "coffee" and "carburettor" share almost none. Any method that compresses those co-occurrence patterns into short vectors puts "coffee" and "tea" close together. There are two families: predict the context with a small neural network (word2vec), or factorize a big table of co-occurrence counts (GloVe). Both end up with similar geometry.

Step 2: word2vec: skip-gram and CBOW

word2vecA family of shallow neural models (skip-gram and CBOW) that learn word vectors by predicting words from their neighbours in a sliding window. (2013) proposed two simple architectures:

• CBOW (continuous bag of words) predicts the current word from its context: given "she poured the ___ into her", predict "coffee".
• Skip-gram predicts the surrounding words given the current word: given "coffee", predict "poured", "into" and "cup".

Each word has a vector, and the prediction is made from dot products between vectors, so the only way to predict contexts well is to give words with similar contexts similar vectors. The task itself is a pretext: nobody needs the predictions. The learned vectors are the product. The paper reports learning high-quality vectors from a 1.6 billion word data set in less than a day.

Predicting one word out of a whole vocabulary needs a softmax over every word, which is slow. The follow-up paper replaced it with negative samplingA training trick for word2vec that turns word prediction into binary classification: score the true context word high and a few randomly drawn words low.: for each real (word, context) pair, also draw k random "noise" words, and train a logistic classifier ( Logistic regression ) to tell the real pair from the fake ones. The authors found k between 5 and 20 useful for small data sets and 2 to 5 enough for large ones, drawing noise words from the unigram distribution raised to the 3/4 power so rare words are sampled a bit more often than their raw frequency.

Here $v_w$ is the vector of the centre word, $v_c$ that of a real context word and $v_{n_i}$ those of the k noise words. Maximizing it pulls real pairs together and pushes random pairs apart.

Step 3: GloVe: factorize the co-occurrence table

GloVeGlobal Vectors: word embeddings trained by fitting dot products of word vectors to the logarithms of global word-word co-occurrence counts. (2014) starts from the full word-word co-occurrence matrix of the corpus. Its authors observe that ratios of co-occurrence probabilities carry meaning: "solid" appears far more often near "ice" than near "steam", "gas" the reverse, while words like "water" and "fashion" relate equally to both, so their ratio is near 1 and cancels out. GloVe trains vectors whose dot products fit the log co-occurrence counts, using only the non-zero entries of the matrix, and describes itself as combining global matrix factorization with local context window methods. In practice word2vec and GloVe vectors behave similarly; pick whichever pretrained set covers your language and domain.

Step 4: Geometry and analogies

Because the vectors are learned from contexts, consistent relationships in the text become consistent offsets in the space. The word2vec paper popularized the test: vector("King") - vector("Man") + vector("Woman") is closest to vector("Queen"). Its test set has 8,869 semantic questions (such as capital cities and their countries) and 10,675 syntactic ones (such as an adjective and its superlative). The analogy is answered by the nearest word to the computed point by cosine similarity ( Vectors and vector spaces ), with the three input words excluded; without that exclusion the nearest word is often just "king".

Treat analogies as a demonstration, not a guarantee. They work for frequent, consistent relationships and fail for many others, and the same arithmetic exposes biases in the training text: offsets between gendered words can also link occupations to genders the way the corpus did.

Going deeper

Code: skip-gram with negative sampling from scratch

import numpy as np
rng = np.random.default_rng(0)

# A synthetic corpus: each person word appears with royal or common, male or female, old or young words
people = {"king": ("royal", "m", "old"), "queen": ("royal", "f", "old"), "prince": ("royal", "m", "young"),
          "princess": ("royal", "f", "young"), "man": ("common", "m", "old"), "woman": ("common", "f", "old"),
          "boy": ("common", "m", "young"), "girl": ("common", "f", "young")}
ctx = {"royal": ["palace", "crown", "throne", "rules"], "common": ["village", "market", "works", "farm"],
       "m": ["he", "his", "him"], "f": ["she", "her", "hers"], "old": ["adult", "wise"], "young": ["child", "plays"]}
sents = []
for _ in range(3000):
    w = str(rng.choice(list(people))); a, g, y = people[w]
    sents.append([w, *(str(rng.choice(ctx[t])) for t in (a, g, y))])
vocab = sorted({t for s in sents for t in s}); ix = {w: i for i, w in enumerate(vocab)}
pairs = [(ix[s[i]], ix[s[j]]) for s in sents for i in range(4) for j in range(4) if i != j]  # skip-gram pairs

# Skip-gram with negative sampling: push real (word, context) pairs together, random pairs apart
V, D, lr, k = len(vocab), 8, 0.05, 3
Win, Wout = rng.normal(0, 0.1, (V, D)), rng.normal(0, 0.1, (V, D))
sig = lambda z: 1 / (1 + np.exp(-z))
for epoch in range(3):
    for n in rng.permutation(len(pairs)):
        w, c = pairs[n]
        for cc, label in [(c, 1)] + [(rng.integers(V), 0) for _ in range(k)]:
            g = sig(Win[w] @ Wout[cc]) - label            # gradient of the logistic loss
            Win[w], Wout[cc] = Win[w] - lr * g * Wout[cc], Wout[cc] - lr * g * Win[w]
print(f"{V} words, {len(pairs):,} training pairs, {D}-dimensional vectors")

E = Win / np.linalg.norm(Win, axis=1, keepdims=True)
def nearest(v, exclude=(), n=3):
    sims = E @ (v / np.linalg.norm(v))
    return [(vocab[i], round(float(sims[i]), 2)) for i in np.argsort(-sims) if vocab[i] not in exclude][:n]
print("nearest to king: ", nearest(E[ix["king"]], ["king"]))
print("nearest to woman:", nearest(E[ix["woman"]], ["woman"]))
for a, b, c in [("king", "man", "woman"), ("prince", "boy", "girl"), ("queen", "woman", "man")]:
    print(f"{a} - {b} + {c} = ", nearest(E[ix[a]] - E[ix[b]] + E[ix[c]], [a, b, c], 1))

26 words, 36,000 training pairs, 8-dimensional vectors
nearest to king:  [('prince', 0.72), ('queen', 0.71), ('man', 0.71)]
nearest to woman: [('queen', 0.73), ('man', 0.72), ('girl', 0.68)]
king - man + woman =  [('queen', 0.98)]
prince - boy + girl =  [('princess', 0.96)]
queen - woman + man =  [('king', 0.98)]

The corpus is synthetic: each person word ("king", "girl", ...) appears in short sentences with words that mark royal or common, male or female, old or young. Nobody tells the model those attributes. Skip-gram with negative sampling (k = 3) learns 8-dimensional vectors from 36,000 (word, context) pairs, and the attributes emerge as directions: "king - man + woman" lands on "queen", "prince - boy + girl" on "princess" and "queen - woman + man" on "king". The nearest neighbours of "king" are "prince", "queen" and "man", each sharing two of its three attributes. Real word2vec is the same loop with a vocabulary of millions, windows over real text, and subsampling of very frequent words.

Limitation: one vector per word

A static embedding gives "bank" one vector, whether the sentence is about money or rivers. The vector ends up as a blend of both senses, close to neither. The follow-up word2vec paper also names two inherent limits of word vectors: they ignore word order, and they cannot represent idiomatic phrases whose meaning is not the sum of the parts, such as "Air Canada". Contextual models fix this by computing a new vector for each occurrence of a word from the whole sentence. That is what attention ( Seq2seq and attention ) and transformers ( Self-attention, step by step ) do, and why modern text embeddings come from them.

Using embeddings

• Pretrained vectors let small labelled data sets benefit from billions of words of unlabelled text: initialize a model's embedding layer with them, or average a sentence's word vectors as a quick document representation.
• Out-of-vocabulary words have no vector. Subword methods ( Text preprocessing and tokenization basics ) such as fastText-style character n-grams can build one from pieces.
• The idea generalizes: anything with contexts can be embedded. Products bought together, songs played in a session and nodes visited in a random walk have all been embedded with the same skip-gram trick ( Two-tower retrieval and embedding-based retrieval ).

Trade-offs

Common mistakes

• Expecting static vectors to separate word senses ("bank", "apple", "python").
• Evaluating analogies without excluding the input words, then concluding the method failed.
• Using vectors trained on news for a specialist domain (medicine, law, code) where key words are rare or mean something else.
• Comparing vectors with Euclidean distance when their lengths vary with word frequency; cosine similarity is the usual choice.
• Forgetting that embeddings learn the biases of their training text and then using them in decisions about people.

In the wild

Check yourself

Recap