Skip to main content

Module 5 — Contextual embeddings and pretrained models

Word2Vec gives apple one vector. The company and the fruit share it, and the geometry reflects some awkward average of the two contexts. Contextual embeddings — ELMo in 2018, BERT and its successors from 2019 on — solve this by producing a different vector for every occurrence, computed from the whole sentence. It is the shift that closed most of the remaining gap with human performance on standardised NLP benchmarks.

Polysemy: one word, several meanings

Consider three sentences drawn from the review corpus:

  • The apple fell off the shelf and bruised badly. — the fruit.
  • Apple's warranty is honoured within 24 hours. — the company.
  • The apple charger died after three months. — a specific product line.

A static Word2Vec vector for apple is a compromise: not really the fruit, not really the company, close to orange and Samsung at the same time. A downstream classifier trying to route reviews to a food team or an electronics team receives the same input for all three sentences.

A contextual embedding gives back three different vectors for the three occurrences of apple. The first sits near pear, bruise, shelf; the second near warranty, service, iPhone; the third near charger, USB, cable. The routing becomes possible.

From ELMo to BERT in one paragraph

ELMo was the first widely used contextual model: a bidirectional LSTM trained to predict the next word from the left context and the previous word from the right context. The representation of each word combined both directions. It worked but stayed sequential and difficult to scale.

BERT replaced the LSTM with a Transformer encoder (course 12) and changed the training objective to masked language modelling: randomly hide 15 % of the tokens in a sentence and ask the model to reconstruct them. The bidirectional attention of the Transformer means every token can look at every other, forward and backward, in parallel — no more sequential bottleneck. Trained on billions of tokens, BERT gave a giant leap on eleven NLP benchmarks with the same architecture, and opened the era of pretrained encoders.

Everything since — RoBERTa, DistilBERT, DeBERTa, ELECTRA — is a refinement of the same recipe: pretrain a Transformer encoder on huge unlabelled text with a self-supervised objective, then fine-tune on the target task (module 6).

Reading a contextual vector

Loading BERT and asking it for a vector takes six lines.

import torch
from transformers import AutoTokenizer, AutoModel

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

encoded = tok("The apple was crisp.", return_tensors="pt")
with torch.no_grad():
out = model(**encoded)

out.last_hidden_state has shape (1, n_tokens, 768). Every token — including the special [CLS] at the start and [SEP] at the end — gets a 768-dimensional vector that already reflects the whole sentence, because every attention layer has mixed information across positions.

The vector for apple in The apple was crisp and the vector for apple in Apple stock jumped will differ noticeably. Cosine similarity between the two is around 0.6 on bert-base-uncased, versus 1.0 for a Word2Vec model. That gap is where the value lives.

The [CLS] illusion for whole-sentence tasks

There is a convenient temptation and a real trap. BERT places a special [CLS] token at the start of every input. During pretraining it was used for a "next sentence prediction" task, so its vector out.last_hidden_state[:, 0, :] is often described as the sentence embedding.

For a fine-tuned classifier, this works: the training procedure adjusts [CLS] to be a good pooled representation for the task. For similarity search out of the box, on a model that has not been fine-tuned for it, [CLS] is a poor sentence vector. Cosine similarity between two paraphrases can be lower than between two unrelated sentences. Many demos silently use it and produce bewildering results.

Sentence-BERT (SBERT) is a family of BERT-style models fine-tuned specifically so that cosine similarity between two sentence vectors matches human notion of similarity. It exposes a single API that returns one vector per sentence, ready to compare.

from sentence_transformers import SentenceTransformer, util

sbert = SentenceTransformer("all-MiniLM-L6-v2") # small and fast

reviews = load_reviews(n=1000) # list of strings
vecs = sbert.encode(reviews, normalize_embeddings=True) # shape (1000, 384)

query = "the battery drains after a few hours"
q = sbert.encode(query, normalize_embeddings=True)

scores = util.cos_sim(q, vecs)[0]
top = scores.topk(5)
for score, idx in zip(top.values, top.indices):
print(f"{score:.3f} {reviews[idx][:80]}")

Two things are worth noticing. First, all-MiniLM-L6-v2 is a 22M-parameter model — small enough to run on CPU at real time — and it beats general-purpose BERT for semantic similarity because the training objective matched the use case. Second, normalize_embeddings=True puts every vector on the unit sphere so a dot product is a cosine similarity: on a million vectors this changes the search from unusable to instant.

Cosine, not Euclidean

The similarity metric matters. Two vectors of very different lengths, in the same direction, are nearly identical in cosine and far in Euclidean. Text embeddings vary in length for reasons that reflect sentence length more than sentence content; cosine cancels that out. Every reasonable semantic search index (FAISS, Milvus, Elasticsearch dense retrieval) offers cosine as the default and normalises internally.

Semantic search on the review corpus

Once every review has a vector, you can ask questions in plain English that TF-IDF could not answer. problems with charging in cold weather retrieves reviews that never contain those exact words but describe the same complaint. This is what unlocks the internal knowledge search of the module 10 project and is at the heart of retrieval-augmented systems in production today.

Fine-tuned for retrieval or not: the vectors are not interchangeable

A bert-base-uncased [CLS] vector and an all-MiniLM-L6-v2 sentence vector do not live in the same space, do not have the same dimension, and cannot be compared. Pick one model per index, and rebuild the index if you change it.

Start with a small SBERT model

all-MiniLM-L6-v2 (22M parameters, 384 dimensions) is often within one or two points of a much larger model on retrieval benchmarks, at a tenth of the memory and latency. For a first index of a few million documents on a laptop, it is the default.

In summary

  • Contextual embeddings give one vector per occurrence, computed from the whole sentence; polysemy stops being a problem the way it was for Word2Vec.
  • BERT trained a Transformer encoder with masked language modelling; every descendant since is a refinement of that recipe.
  • Using the raw [CLS] vector for similarity search is a common mistake; Sentence-BERT models are fine-tuned so that cosine similarity between sentence vectors matches human judgement.
  • Cosine similarity on unit-normalised vectors is the right metric for text; a small SBERT model plus a proper index gives semantic search on a laptop.

Next module: taking a pretrained encoder and fine-tuning it to classify the same reviews we vectorised here.