Embedding Space — Module 11.1

How This Lesson Fits the Module & Volume

The previous lecture defined the embedding table. Embedding space is the geometry of those vectors (and, relatedly, of contextual hidden states): distances, angles, clusters, and linear structure. Vol. 09 used this intuition for Word2Vec analogies; Vol. 05’s PCA / t-SNE tools still help visualize it.

In LMs, the same geometry underpins nearest-neighbor diagnostics, retrieval, and why the LM head—a linear map into logits—can be read as scoring directions in space.

Learning Objectives

By the end of this lesson, students should be able to:

Describe embedding space as R^d populated by token (or state) vectors.
Compute cosine similarity and L2 distance between embedding rows in PyTorch.
Interpret clusters and linear analogies as useful—but imperfect—structure.
Contrast input embedding space with contextual hidden-state space.
Relate the LM head to scoring alignment between a state and vocab directions.
Use geometry as a debugging lens before moving on to hidden states and logits.

Definition

Embedding space is the d-dimensional vector space in which tokens (or hidden states) are represented. Similarity is typically measured by cosine similarity (angle) or Euclidean distance; learned training objectives shape which points land near each other.

Metrics That Matter

Measure	Formula intuition	Use when
Cosine similarity	Alignment of directions	Magnitude is less meaningful
Dot product	Alignment × magnitudes	Matches attention / LM-head scores
L2 distance	Straight-line separation	Clustering, some retrieval stacks

Input Embedding Space

|V| points (one per type).
Static until fine-tuned.
Good for type-level neighbors.

Hidden-State Space

One vector per position.
Context-sensitive.
Used for probing / features.

LM-Head View

Each vocab ID has a direction.
Logit ≈ state · direction (+ bias).
Next-token pick = nearest / highest score.

Code: Nearest Neighbors in GPT-2 Embeddings

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "gpt2"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name)
E = model.transformer.wte.weight.detach()          # (|V|, d)
E_n = F.normalize(E, dim=-1)

def neighbors(word, k=8):
    tid = tok.encode(word, add_special_tokens=False)
    assert len(tid) == 1, "pick a single-token word for this demo"
    q = E_n[tid[0]]
    scores = E_n @ q
    top = torch.topk(scores, k + 1)
    for s, i in zip(top.values.tolist(), top.indices.tolist()):
        piece = tok.decode([i])
        if i == tid[0]:
            continue
        print(f"{s:.3f}  {piece!r}")

neighbors(" cat")   # GPT-2 often has leading-space tokens

Structure and Its Limits

Word2Vec-style analogies (king − man + woman ≈ queen) popularized linear structure. LM embedding spaces show related clusters (digits, languages, code punctuation) but are shaped by subword segmentation and next-token statistics—not a clean semantic ontology. Treat neighbor lists as diagnostics, not proof of understanding.

Why Geometry Helps

Debug tokenizer/model mismatch (neighbors look random).
Intuition for retrieval and RAG embeddings.
Explains LM head as directional scoring.

Pitfalls

Cosine on subword pieces can mislead.
2D projections (t-SNE) distort distances.
Input-space neighbors ≠ contextual synonyms.

Common Misconception

“If two words are synonyms, their embedding rows must be nearest neighbors.” Synonymy is contextual, and many synonyms tokenize into multiple pieces. Even for single-token types, training pressures (frequency, syntax) can place vectors near distributional co-travelers rather than thesaurus synonyms. Always verify with the actual tokenizer pieces.

Bridge to the Rest of Module 11.1

From here the curriculum moves inside the stack: hidden states evolve through layers; logits and softmax turn final states into the next-token distribution you met earlier. Embedding space is the first geometric chapter—not the last.

Knowledge Check

Short Answer: What is embedding space? Answer: The d-dimensional vector space where token (or state) vectors live.
True/False: Cosine similarity ignores vector magnitude (after normalization). Answer: True—it depends on angle / normalized direction.
Multiple Choice: LM-head logits are closest to: (a) edit distance, (b) dot products with vocab directions, (c) TF-IDF. Answer: (b).
Short Answer: Name one difference between input embedding space and hidden-state space. Answer: Input space is type-level/static per ID; hidden states are contextual per position.
True/False: t-SNE plots preserve all high-d distances faithfully. Answer: False.
Multiple Choice: A useful debugging signal for a mismatched tokenizer is: (a) perfect BLEU, (b) nearest neighbors look like noise / wrong language pieces, (c) lower VRAM. Answer: (b).
Short Answer: Why might neighbors("cat") fail for GPT-2? Answer: The type may be stored as a leading-space token like " cat", or split into multiple pieces.
Short Answer: What Vol. 05 tools help visualize embedding space? Answer: PCA, t-SNE, or UMAP (any of these).
Multiple Choice: Distributional neighbors tend to reflect: (a) only WordNet synonyms, (b) patterns from training co-occurrence / next-token stats, (c) random init forever. Answer: (b).
True/False: Geometry alone proves an LM “understands” a concept. Answer: False.

Key Takeaways

Embedding space is R^d with similarity structure shaped by training.
Cosine / dot product / L2 answer different geometric questions.
Input embeddings and contextual hidden states are related but not identical spaces.
The LM head scores directions in this geometry to produce next-token logits.
Next: Hidden State—vectors after attention mixes context.

Trainer’s Guide

Hands-on idea: Have students find nearest neighbors for a punctuation piece, a digit, and a code keyword; discuss clusters.

Discussion prompt: When should a RAG system embed with input token averages vs a dedicated sentence embedding model?

Recap: Embedding space is the geometry of token vectors—use similarity carefully, then move into contextual hidden states. Continue with Hidden State.