Contextual Embeddings: ELMo & the Path to Transformers

A word's meaning depends on context, they said. So does my performance review, but nobody built ELMo for that.

Lexica, Context-Dependent AI Agent

The Polysemy Problem

Word2Vec, GloVe, and FastText share a fundamental limitation: each word gets exactly one vector, regardless of context. But language is deeply contextual.

Consider the word "bank":

• "I deposited money at the bank." → financial institution
• "We sat on the river bank." → edge of a river
• "Don't bank on it." → to rely on

With Word2Vec, all three uses map to the same vector: a compromise that captures none of the meanings well. This is called the polysemy problem.

Figure 1.4.2: Static vs. contextual embeddings. Word2Vec assigns one fixed vector per word regardless of context, while contextual models produce different vectors for each usage.

ELMo: Embeddings from Language Models (2018)

ELMo (Peters et al., 2018) was the first widely successful contextual embedding model. The key idea: run the entire sentence through a deep bidirectional LSTM, and use the hidden states as word representations. Since the LSTM has seen the whole sentence, each word's representation is influenced by its context.

How it works:

1. Train a bidirectional language model (forward LSTM + backward LSTM) on a large corpus
2. For each word in a sentence, extract hidden states from all layers
3. The ELMo embedding is a learned weighted combination of all layer representations

Figure 1.4.3: The ELMo computation graph. Input tokens pass through a character embedding layer, then through forward and backward LSTMs. The final ELMo vector for each token is a learned weighted sum across all layers.

The breakthrough: "bank" in "river bank" now gets a different vector than "bank" in "bank account", because the LSTM hidden states are conditioned on the entire sentence.

Why Different Layers Capture Different Information

A remarkable finding from the ELMo paper: different layers of the LSTM capture different types of linguistic information:

• Layer 0 (token embeddings): Captures basic word identity and morphological features, similar to Word2Vec. "Running" is close to "runner" and "runs."
• Layer 1 (first LSTM): Captures syntactic information: part of speech, grammatical role. The model has learned whether a word is a noun, verb, or adjective from context.
• Layer 2 (second LSTM): Captures semantic information: word sense disambiguation, topical context. This is the layer that distinguishes "river bank" from "bank account."

Figure 1.4.5: The information hierarchy across ELMo layers. Lower layers capture morphology and word identity, middle layers encode syntax, and upper layers represent semantics and word sense.

This is why ELMo uses a weighted combination of all layers (the α weights in the diagram above): different downstream tasks benefit from different layers. A POS tagger might weight Layer 1 heavily, while a sentiment classifier might rely more on Layer 2. The weights are learned during fine-tuning for each specific task.

Contextual Embeddings in Code

While ELMo itself is rarely used directly today (BERT and transformers have superseded it), we can demonstrate the concept of contextual embeddings using Hugging Face Transformers, which makes it easy to extract hidden states from any model: Code Fragment 1.4.1 below puts this into practice.

# Demonstrating contextual embeddings: same word, different vectors
from transformers import AutoTokenizer, AutoModel
import torch
from scipy.spatial.distance import cosine

# Load a BERT model (the modern successor to ELMo)
# NOTE: First run will download bert-base-uncased (~420 MB).
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

def get_word_embedding(sentence, word):
 """Extract the contextual embedding for a specific word in a sentence."""
 inputs = tokenizer(sentence, return_tensors="pt")
 with torch.no_grad():
 outputs = model(**inputs)

 # Find the token position for our word
 # WARNING: This assumes the target word is a single token. If BERT's tokenizer
 # splits the word into subwords (e.g., "deposit" -> "dep" + "##osit"), this
 # lookup will fail. For production code, you would need to handle multi-token
 # words by averaging their subword embeddings.
 tokens = tokenizer.tokenize(sentence)
 word_idx = tokens.index(word) + 1 # +1 for [CLS] token

 # Return the hidden state at that position
 return outputs.last_hidden_state[0, word_idx].numpy()

# "bank" in two different contexts
bank_river = get_word_embedding("I sat by the river bank", "bank")
bank_money = get_word_embedding("I went to the bank to deposit money", "bank")

# Measure how different the two "bank" vectors are
distance = cosine(bank_river, bank_money)
print(f"Cosine distance between 'bank' in different contexts: {distance:.3f}")
# Output: ~0.35: substantially different vectors for the same word!

# Compare: "bank" (river) is closer to "shore" than to "bank" (money)
shore = get_word_embedding("We walked along the shore", "shore")
print(f"Distance bank(river) to shore: {cosine(bank_river, shore):.3f}")
print(f"Distance bank(river) to bank(money): {cosine(bank_river, bank_money):.3f}")
# bank(river) is CLOSER to "shore" than to bank(money)!
# This is exactly what contextual embeddings solve.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = model.encode([
 "I sat by the river bank",
 "I went to the bank to deposit money",
 "We walked along the shore"
])
# embeddings.shape: (3, 384), ready for cosine similarity

From ELMo to Transformers: What Changed

Let us compare the approaches side by side to see the progression clearly:

The key limitation of ELMo was its reliance on LSTMs, which process text sequentially (one word at a time). This made training slow and limited the model's ability to capture very long-range dependencies. The Transformer architecture (Chapter 4) solves this by processing all words simultaneously using self-attention.

Summary: The Representation Journey

This module traced the evolution of how we represent text for machines. Each step solved a problem that the previous approach could not handle:

What is Next

In Chapter 2, we will explore tokenization, the critical first step that determines how text is broken into pieces before being fed to any model. You will learn the BPE algorithm that powers GPT-4 and Llama, and understand why tokenizer choice affects everything from model quality to API cost.

In Chapter 3, we will dive deep into attention, the mechanism that solved the sequential bottleneck of RNNs and enabled the Transformer revolution.

And in Chapter 4, we will build a complete Transformer from scratch in PyTorch, the architecture that underpins every model we will work with for the rest of the book.

Exercises & Self-Check Questions

Conceptual Questions

1. Representation evolution: In your own words, explain why the transition from sparse vectors (BoW) to dense vectors (Word2Vec) was such a big deal. What specific problems did it solve, and what new capabilities did it unlock?
2. The distributional hypothesis: The phrase "you shall know a word by the company it keeps" is the foundation of Word2Vec. Can you think of cases where this assumption breaks down? (Hint: think about antonyms. "Hot" and "cold" appear in very similar contexts...)
3. Static vs. contextual: Give three sentences where the word "play" means different things. Explain why Word2Vec would struggle with these but ELMo would handle them well.
4. Why pre-train? ELMo was pre-trained on a large corpus, then used for specific tasks. Why is this better than training a model from scratch for each task? What does the pre-training capture that task-specific training would miss?
5. Trade-offs: A colleague argues "TF-IDF is obsolete; just use embeddings for everything." Give two scenarios where TF-IDF would actually be the better choice, and explain why.

Coding Exercises

1. Preprocessing exploration: Take a paragraph from a news article. Run it through the preprocessing pipeline from Section 1.2. Then experiment: what happens if you do not remove stop words? What if you use stemming instead of lemmatization? How do the resulting BoW vectors differ?
2. Analogy hunting: Using the pre-trained word2vec-google-news-300 vectors, find 5 analogies that work well (beyond king/queen) and 3 that fail. Can you explain why the failures happen?
3. Similarity exploration: Pick 20 words from three different categories (e.g., sports, food, technology). Compute all pairwise cosine similarities. Do words within a category have higher similarity than cross-category pairs? Visualize the similarity matrix as a heatmap.
4. Word2Vec from scratch (challenge): Implement the Skip-gram model with negative sampling in pure PyTorch (no Gensim). Train it on a small text corpus and verify that similar words end up with similar vectors. Compare your results with Gensim's output.
   
   Scaffolding hint: Your implementation will need four key components: (1) a SkipGramDataset class that generates (center, context, negative) tuples from your corpus; (2) two nn.Embedding layers, one for center words and one for context words; (3) a negative sampling loss function that maximizes the dot product for positive pairs and minimizes it for negative pairs (see the formula in Section 1.3); and (4) a training loop that iterates over batches, computes the loss, and calls optimizer.step(). Start with a small vocabulary (a few hundred words) and short embedding dimension (50) to debug before scaling up.

Further Reading

pip install gensim matplotlib scikit-learn transformers torch numpy

from gensim.models import Word2Vec

# Sample corpus (in practice, use a larger dataset)
sentences = [
 ["the", "king", "rules", "the", "kingdom"],
 ["the", "queen", "rules", "the", "kingdom"],
 ["a", "prince", "is", "the", "son", "of", "a", "king"],
 ["a", "princess", "is", "the", "daughter", "of", "a", "king"],
 ["the", "man", "works", "in", "the", "city"],
 ["the", "woman", "works", "in", "the", "city"],
 ["a", "boy", "plays", "in", "the", "park"],
 ["a", "girl", "plays", "in", "the", "park"],
 ["the", "dog", "runs", "in", "the", "park"],
 ["the", "cat", "sleeps", "on", "the", "couch"],
] * 100 # Repeat for better training

model = Word2Vec(sentences, vector_size=50, window=3,
 min_count=1, sg=1, epochs=50)

# Query nearest neighbors
for word in ["king", "queen", "man"]:
 neighbors = model.wv.most_similar(word, topn=3)
 print(f"{word}: {[(w, f'{s:.2f}') for w, s in neighbors]}")

import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import numpy as np

words = list(model.wv.key_to_index.keys())
vectors = np.array([model.wv[w] for w in words])

tsne = TSNE(n_components=2, random_state=42, perplexity=min(5, len(words) - 1))
coords = tsne.fit_transform(vectors)

fig, ax = plt.subplots(figsize=(10, 8))
ax.scatter(coords[:, 0], coords[:, 1], s=40, alpha=0.6)
for i, word in enumerate(words):
 ax.annotate(word, (coords[i, 0], coords[i, 1]),
 fontsize=9, ha="center", va="bottom")
ax.set_title("Word2Vec Embeddings (t-SNE projection)")
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("word2vec_tsne.png", dpi=150)
plt.show()

from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
bert = AutoModel.from_pretrained("bert-base-uncased")

sentences = [
 "I deposited money at the bank",
 "We sat on the river bank",
 "The bank approved the loan",
 "Fish swim near the bank of the stream",
]

def get_word_embedding(sentence, target_word):
 """Extract BERT embedding for a target word in context."""
 inputs = tokenizer(sentence, return_tensors="pt")
 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
 with torch.no_grad():
 outputs = bert(**inputs)
 # Find the target word's token index
 idx = next(i for i, t in enumerate(tokens) if target_word in t)
 return outputs.last_hidden_state[0, idx].numpy()

# Get "bank" embeddings in each context
embeddings = [get_word_embedding(s, "bank") for s in sentences]

# Compute pairwise cosine similarity
from numpy.linalg import norm
print("Cosine similarity between 'bank' in different contexts:")
for i in range(len(sentences)):
 for j in range(i + 1, len(sentences)):
 cos_sim = np.dot(embeddings[i], embeddings[j]) / (
 norm(embeddings[i]) * norm(embeddings[j]))
 print(f" [{i}] vs [{j}]: {cos_sim:.3f}")
 print(f" '{sentences[i]}' vs '{sentences[j]}'")

from numpy.linalg import norm

n = len(embeddings)
sim_matrix = np.zeros((n, n))
for i in range(n):
 for j in range(n):
 sim_matrix[i, j] = np.dot(embeddings[i], embeddings[j]) / (
 norm(embeddings[i]) * norm(embeddings[j]))

fig, ax = plt.subplots(figsize=(6, 5))
im = ax.imshow(sim_matrix, cmap="RdYlGn", vmin=0.7, vmax=1.0)
short_labels = ["money at bank", "river bank", "bank approved", "bank of stream"]
ax.set_xticks(range(n))
ax.set_xticklabels(short_labels, rotation=45, ha="right", fontsize=9)
ax.set_yticks(range(n))
ax.set_yticklabels(short_labels, fontsize=9)
for i in range(n):
 for j in range(n):
 ax.text(j, i, f"{sim_matrix[i,j]:.2f}", ha="center", va="center", fontsize=10)
ax.set_title("Contextual Similarity of 'bank' (BERT)")
plt.colorbar(im, ax=ax, shrink=0.8)
plt.tight_layout()
plt.savefig("contextual_bank.png", dpi=150)
plt.show()