This week we go deep into embeddings — how discrete token IDs become dense, meaningful vectors, how those vectors encode semantic relationships, and how to work with them in practice.
Embeddings and why we need them see section on categorical data and encoding
words in similar context have similar meanings
we need embeddings which take care of context
embeddings optimized for specific task. See why and see notes for UG class here and here and notes and notes
Lecture Prep: Fundamentals of Embeddings and Large Language Models
However, simple methods like assigning a unique integer to each word or using one-hot encoding fall short because they fail to capture meaning.
1 and “orange” is assigned 2, a mathematical model treats them as completely independent units. The numerical distance between 1 and 2 holds no intrinsic semantic relationship, meaning the model cannot infer that both are types of fruit.A crucial point for students to grasp is that embeddings are not static or absolute representations; they are tied explicitly to a specific objective and are trained end-to-end.
Context & Task Definition: When a representation is optimized for a particular task, the dimensions of the embedding learn to reflect the properties most useful for solving that task. The objective dictates the underlying geometry of the vector space.
Example: Word Prediction: If the target task is next-word prediction (the foundational objective of LLMs), the trained embeddings will manifest as highly specific numbers optimized exclusively to facilitate that prediction loop.
End-to-End Optimization: Rather than manually assigning features, initial embedding weights are usually randomized. Through the process of backpropagation, the gradients flow all the way back to the input embedding layer. The network iteratively tunes these vector coordinates until the spatial layout minimizes the error rate of the target task.
For highly visual and intuitive breakdowns of these exact concepts, the following video lessons from the 3Blue1Brown channel are exceptional tools for both lecture planning and student reference:
After tokenisation we have a sequence of integers, e.g. [1045, 2293, 4671]. A neural network cannot directly compute on integers. We need to map each integer to a continuous vector.
An embedding layer is simply a lookup table:
E : V → R^d
where V is the vocabulary size (e.g. 50,257) and d is the embedding dimension (e.g. 768 or 4096). Each token ID is mapped to a vector of d floats.
In PyTorch:
import torch
import torch.nn as nn
vocab_size = 50257
d_model = 768
embedding = nn.Embedding(vocab_size, d_model)
token_ids = torch.tensor([1045, 2293, 4671])
vectors = embedding(token_ids) # shape: (3, 768)
These vectors are learned during training — they start random and gradually encode useful structure.
“You shall know a word by the company it keeps.” — J. R. Firth (1957)
Words that appear in similar contexts tend to have similar meanings. By training on next-token prediction, the model is forced to give similar representations to tokens that behave similarly in text.
Mikolov et al. (2013) introduced Word2Vec — a simple two-layer neural network trained with two objectives:
The result: embeddings that capture analogy relationships:
vec("king") - vec("man") + vec("woman") ≈ vec("queen")
The embedding encodes dimensions corresponding loosely to:
These are not explicitly programmed — they emerge from co-occurrence statistics.
Word2Vec gives a static embedding: “bank” has one vector regardless of context (river bank vs financial bank). This is a major limitation.
Modern LLMs produce contextual embeddings: the representation of a token depends on all surrounding tokens. The same word gets a different vector in different sentences.
This is a direct consequence of the attention mechanism (covered next week).
For tasks like semantic search or clustering, we need a single vector representing an entire sentence or document.
Average the token embeddings:
sentence_embedding = token_embeddings.mean(dim=0)
Use the embedding of the special [CLS] token, which is designed to summarise the full input.
sentence-transformers library: fine-tuned models optimised for semantic similarity (e.g. all-MiniLM-L6-v2, all-mpnet-base-v2).from sentence_transformers import SentenceTransformer
model = SentenceTransformer("all-MiniLM-L6-v2")
sentences = ["The cat sat on the mat.", "A feline rested on a rug."]
embeddings = model.encode(sentences) # shape: (2, 384)
Given two embedding vectors u and v, the most common similarity metric is cosine similarity:
cos_sim(u, v) = (u · v) / (||u|| * ||v||)
import torch
import torch.nn.functional as F
sim = F.cosine_similarity(u.unsqueeze(0), v.unsqueeze(0))
High-dimensional embeddings (d = 768, 4096) cannot be visualised directly. We use dimensionality reduction:
Projects to 2D by finding the axes of maximum variance. Fast, linear, interpretable.
Non-linear projection. Preserves local structure — similar vectors cluster together. Good for visualising cluster structure.
Faster than t-SNE, better at preserving global structure. Widely used in modern practice.
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
tsne = TSNE(n_components=2, random_state=42)
reduced = tsne.fit_transform(embeddings)
plt.scatter(reduced[:, 0], reduced[:, 1])
for i, word in enumerate(words):
plt.annotate(word, reduced[i])
plt.show()
The transformer has no built-in notion of order — the same set of tokens in different orders would produce the same result without position information. We add positional encodings to the embeddings.
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
Simply learn a separate embedding for each position. Works well in practice but does not generalise to longer sequences.
Encodes position as a rotation of the query and key vectors. Better generalisation to longer contexts. This is the modern standard.
See practicals/week03_practical.py:
king - man + woman ≈ queen.