Embeddings & Representation Learning
What an embedding actually is mechanically, why training makes embedding spaces semantically structured, and how contrastive learning builds them on purpose
Embeddings & Representation Learning
TL;DR
An embedding is a learned vector produced by a lookup table (nn.Embedding), trained so that vectors for related things end up close together — not because anyone told the model what's related, but because the training objective mathematically rewards it. Contrastive learning is the training approach that builds this deliberately: pull matching pairs together, push non-matching pairs apart, and a coherent geometric structure falls out.
| Property | Value |
|---|---|
| Level | Intermediate |
| Reading time | ~18 minutes |
| Prerequisites | Regularization & Generalization |
| You will understand | What an embedding is mechanically, why embedding spaces end up semantically structured, and how contrastive learning builds that structure on purpose |
What an Embedding Actually Is
Mechanically, an embedding is nothing exotic — it's a row in a lookup table.
import torch.nn as nn
vocab_size, embedding_dim = 50_000, 768
embedding_table = nn.Embedding(vocab_size, embedding_dim)
token_id = torch.tensor([42])
vector = embedding_table(token_id) # shape: (1, 768) — just a row lookupnn.Embedding is a matrix of shape (vocab_size, embedding_dim), and indexing into it with a token ID returns that row. Every one of those rows starts as random noise. What makes the final table useful is entirely a product of training: gradients flow back into whichever rows were looked up during a batch, nudging them in whatever direction reduces the loss.
There's no embedding "algorithm" beyond the lookup. All the interesting behavior — words clustering by meaning, similar images landing near each other — comes from what the training objective rewards, not from anything built into nn.Embedding itself. Change the objective and the same architecture produces a completely different, differently-structured space.
Why the Space Ends Up Structured
Consider a simple objective: predict a word from the words around it (the core idea behind classic word2vec-style training). To do that well, the model has no choice but to place words that appear in similar contexts — "cat" and "dog" both often follow "I have a" and precede "ran away" — at similar points in the vector space, because a shared surrounding context is exactly what the training signal rewards putting close together.
Training signal: "predict the missing word from its context"
"The ___ chased the mouse." → cat, dog, fox all plausible
"I walked my ___ this morning." → dog, cat (less so), fox (no)
Words that fit similar contexts get pushed toward similar vectors —
not because they're told to, but because doing so reduces the loss
on many training examples at once.This generalizes far beyond word prediction. Any objective that rewards a model for treating semantically related inputs similarly — next-token prediction, masked reconstruction, or an explicit contrastive objective — will tend to produce a geometrically structured space as a side effect of minimizing its loss.
Comparing Embeddings: Cosine Similarity and Dot Product
Once you have vectors, "how similar are these two things" becomes a geometry question.
import torch.nn.functional as F
cosine_sim = F.cosine_similarity(vec_a, vec_b, dim=-1)
# cos(θ) = (a · b) / (‖a‖ ‖b‖) — ranges from -1 (opposite) to 1 (identical direction)
dot_product = (vec_a * vec_b).sum(dim=-1)
# a · b — also grows with vector magnitude, not just direction| Metric | Sensitive to magnitude? | Typical use |
|---|---|---|
| Cosine similarity | No — normalizes it out | Comparing direction, the most common choice for semantic similarity |
| Dot product | Yes | Used directly when the embedding model was specifically trained so magnitude also carries signal (some retrieval-tuned models) |
Most embedding models are trained with cosine similarity (or a closely related normalized metric) as the training-time objective itself — which means at inference time, using the same metric the model was optimized for isn't a stylistic choice, it's what the vectors were actually shaped to be compared with.
Contrastive Learning: Building the Structure on Purpose
Word-prediction-style objectives produce useful structure as a side effect. Contrastive learning builds it directly: explicitly train the model so that known-matching pairs end up close, and known-non-matching pairs end up far apart.
One contrastive training step
A simplified InfoNCE-style contrastive loss, for one positive pair against a batch of negatives:
def contrastive_loss(anchor, positive, negatives, temperature=0.05):
pos_sim = F.cosine_similarity(anchor, positive, dim=-1) / temperature
neg_sims = F.cosine_similarity(anchor.unsqueeze(1), negatives, dim=-1) / temperature
logits = torch.cat([pos_sim.unsqueeze(1), neg_sims], dim=1)
labels = torch.zeros(logits.shape[0], dtype=torch.long) # index 0 = the positive
return F.cross_entropy(logits, labels)This is literally a classification problem in disguise: "given this anchor, which of these candidates (one real positive, several negatives) is the match?" — and minimizing that cross-entropy loss is exactly what pulls the positive's vector toward the anchor's and pushes every negative's vector away.
The temperature parameter controls how sharply the loss punishes near-misses. A low temperature makes the softmax over similarities much sharper, so the loss aggressively penalizes any negative that's even moderately close to the anchor — useful when you want fine-grained separation, but can make training unstable if set too low. This is the same softmax-temperature idea used in Knowledge Distillation Mechanics, applied to a different loss.
This is the training approach behind most modern embedding models — the ones used practically for retrieval in Embeddings Explained and loaded via sentence-transformers in practice were trained exactly this way, on millions of positive/negative pairs.
The Geometry This Produces: Vector Arithmetic
A well-trained embedding space doesn't just cluster similar things — it can encode relationships as consistent directions:
embed("king") - embed("man") + embed("woman") ≈ embed("queen")This works because "royalty" and "gender" end up represented as roughly consistent directions across many word pairs, not because the model was ever told about kings or queens specifically. If "man → king" and "woman → queen" both correspond to adding approximately the same "royalty" offset, then vector arithmetic that combines those offsets lands near the expected word. It's a striking illustration of what "semantically structured" concretely means: not just proximity, but consistent, composable directions.
This exact arithmetic doesn't hold perfectly for every relationship in every embedding space — it's a famous, clean illustration from classic word embeddings, and the degree to which it holds varies by model and relationship type. The underlying point generalizes better than the specific trick: training pushes related concepts into geometrically consistent relative positions.
Concept Checks
Check yourself
Why does an embedding table trained only to predict a word from its context end up placing semantically related words near each other, when it was never told which words are related?
Because words appearing in similar contexts get similar gradient updates during training — if "cat" and "dog" both frequently appear after "I have a," the loss is reduced by making their vectors respond similarly to that context, which geometrically means pushing them toward similar positions. The semantic clustering is an emergent side effect of minimizing the prediction loss, not something explicitly encoded.
Why is cosine similarity usually the right metric to compare two embeddings, rather than raw Euclidean distance?
Because most embedding models are trained with a cosine-similarity-based (or closely related normalized) objective, so the vectors are shaped specifically to be compared by direction rather than magnitude. Using the same metric the model was optimized against at training time is what makes the comparison meaningful — a different metric wasn't what the training process was actually pushing the vectors to respect.
In the contrastive loss shown, what does lowering the temperature parameter actually change about training?
It sharpens the softmax over similarity scores, so the loss punishes negatives that are even moderately similar to the anchor much more aggressively than a higher temperature would. This pushes the model toward finer-grained separation between the positive and near-miss negatives, at the cost of potentially less stable training if set too aggressively low.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
nn.Embedding | A learned lookup table — a row per token/item, trained via ordinary gradient descent |
| Emergent structure | Semantic clustering arises from the training objective rewarding it, not from anything built into the layer |
| Cosine similarity | Compares vector direction, normalized for magnitude — matches how most embedding models are trained |
| Contrastive learning | Explicitly pulls known-positive pairs together and pushes negatives apart |
| InfoNCE-style loss | Frames "which candidate matches the anchor" as a classification problem over positives and negatives |
| Temperature | Controls how sharply the contrastive loss punishes near-miss negatives |
| Vector arithmetic | Consistent relationship directions (e.g. royalty, gender) emerging from training, illustrated by king − man + woman ≈ queen |
Next
With a solid mental model of how a representation space gets built, the next page covers adapting an already-trained model to a new task: Fine-Tuning & Transfer Learning.
Regularization & Generalization
Overfitting, the train/validation gap as the central diagnostic, and the mechanics of dropout, weight decay, and normalization
Fine-Tuning & Transfer Learning
Why transfer learning works, how layer freezing works mechanically, and the full derivation of LoRA's low-rank update