Chapter 1 · Section 1.10 · 55 min

The embedding map

You started this chapter with a single question, how to hand the word "cat" to a computer. The answer, a learned vector whose geometry carries meaning, turns out to be the thread that runs through essentially every modern AI system. Language models, image generators, voice assistants, search engines, recommenders, robots that plan, and agents that remember are all, at their core, machines that make embeddings, transform embeddings, compare embeddings, and decode embeddings.

This last section is a guided tour of that landscape. Nothing here is meant to be mastered yet. Each stop is a part of the book still ahead. The goal is that when you arrive at each one, you'll recognize it as an old friend wearing new clothes, and that the small, exact pieces you build here (the output head of a language model, a bridge between two spaces, and a demonstration of collapse) are already the real thing.

Stop 1. Everything enters as an embedding

Every model begins by turning its input into vectors, and you've now seen all the main ways.

Tables for discrete things such as tokens, users, and items (§1.2 and §1.7).
Projections for continuous chunks such as image patches, audio frames, and video tubelets (§1.8).

One more ingredient is added to almost all of them, an embedding that says where something is. You saw in §1.9 that a bag of word vectors cannot tell "dog bites man" from "man bites dog". The attention layers at the heart of a Transformer have the same blindness, because they treat their input as a set of vectors rather than a sequence. The fix is to add a position vector to each token vector, either a learned row of a second table (one row per position, the same lookup as §1.2) or a fixed pattern of sines and cosines at many frequencies, like the hands of a clock turning at different speeds. The original Transformer used the sinusoids, and today's language models instead rotate their vectors by position-dependent angles. §5.4 builds the sinusoidal and learned codes and proves that shifting a sinusoidal code by a fixed offset is a rotation, and §5.5 builds the rotary scheme. The same sinusoidal trick also tells diffusion models how much noise is in their input, by embedding the noise level as a vector of sines and cosines.

Stop 2. A language model is an embedding refinery

Inside a language model, each token's vector passes up through a stack of layers. Every layer lets the vectors exchange information with the other positions and adds an update to each one (§5.1 calls this running sum the residual stream). The vector for "bank" starts as the static table row from §1.2 and, layer by layer, becomes a description of this bank in this sentence. That is the fix for the first cracks in §1.9.

At the top, the model must choose the next token. It does that with the operation you know best, a dot product of the last position's final vector $\mathbf{h} \in \mathbb{R}^d$ (its final hidden state) with every token's output embedding,

p(\text{next} = w \mid \text{context}) = \operatorname{softmax}\!\big(\mathbf{E}_{\text{out}}\,\mathbf{h}\big)_w , \qquad \mathbf{E}_{\text{out}} \in \mathbb{R}^{V\times d}.

In plain words, the next token is the one whose output vector points most in the direction of the final hidden state. With weight tying (§1.2), $\mathbf{E}_{\text{out}} = \mathbf{E}$ , the very same table that embedded the input. Seen this way, a language model is the tied bigram model of §1.2 with a deep stack of layers inserted between the lookup and the output. So an LLM reads by looking up embeddings, thinks by refining them, and speaks by finding which embedding its final vector points toward. Researchers even apply this output head to intermediate layers to watch a prediction form as the vector travels up, a technique called the logit lens. You'll build this whole pipeline when you build a GPT from scratch.

Core

From a hidden state to the next token

Implement next_token(h, E, k, temperature=1.0) returning the indices and probabilities of the k most likely next tokens, given a final hidden state h of shape (d,) and a tied embedding table E of shape (V, d). Use a stable softmax over E @ h / temperature, and return (top_ids, top_probs) sorted by probability. This is the last step of every language model forward pass, and temperature is the knob that §6.4 turns into sampling strategies.

import numpy as np

def next_token(h, E, k, temperature=1.0):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Stop 3. Bridges between spaces

Different encoders produce different spaces. A striking fact of modern AI is how often a simple map, sometimes just a linear one, can connect them.

Vision-language models. LLaVA (Liu et al., 2023) takes the patch embeddings from a CLIP image encoder and passes them through a small projector (a single linear layer in the first version, a two-layer MLP in LLaVA-1.5) into the token embedding space of a language model. In LLaVA-1.5, a $336\times336$ image becomes $24\times24 = 576$ patch vectors of width 1,024, each projected to the 4,096-dimensional token space of its 7B language model. To the LLM, an image is just 576 extra "words" at the start of the prompt. Other designs let the language model look at image features through cross-attention (Flamingo), or first compress the image into a few dozen vectors with a small set of learned queries (BLIP-2 uses 32).
Aligning spaces across languages. Mikolov, Le and Sutskever (2013) found that word2vec spaces trained separately on English and Spanish can be aligned with a linear map learned from a small dictionary, because the shape of meaning is similar across languages. Xing et al. (2015) showed that constraining the map to be orthogonal (a rotation, which preserves all distances and cosines) works better, and Conneau et al. (2018) learned such rotations with almost no dictionary at all.

Here's that idea in miniature. We create a "space A" and a "space B" that are secretly a rotated, noisy copy of each other, and recover the rotation from a few matched pairs.

Python · cell 1

∂ Why the best rotation is U V-transpose

We want the orthogonal $\mathbf{R}$ (with $\mathbf{R}^\top\mathbf{R} = \mathbf{I}$ ) minimizing $\lVert\mathbf{X}\mathbf{R} - \mathbf{Y}\rVert_F^2$ for matched rows $\mathbf{X}, \mathbf{Y} \in \mathbb{R}^{n\times d}$ . Expand with $\lVert\mathbf{M}\rVert_F^2 = \operatorname{tr}(\mathbf{M}^\top\mathbf{M})$ ,

\lVert\mathbf{X}\mathbf{R} - \mathbf{Y}\rVert_F^2 = \operatorname{tr}(\mathbf{R}^\top\mathbf{X}^\top\mathbf{X}\mathbf{R}) - 2\operatorname{tr}(\mathbf{R}^\top\mathbf{X}^\top\mathbf{Y}) + \operatorname{tr}(\mathbf{Y}^\top\mathbf{Y}).

The first term equals $\operatorname{tr}(\mathbf{X}^\top\mathbf{X}\mathbf{R}\mathbf{R}^\top) = \operatorname{tr}(\mathbf{X}^\top\mathbf{X})$ because the trace is cyclic and $\mathbf{R}\mathbf{R}^\top = \mathbf{I}$ , so it doesn't depend on $\mathbf{R}$ . Minimizing the error is therefore maximizing $\operatorname{tr}(\mathbf{R}^\top\mathbf{X}^\top\mathbf{Y})$ . Take the SVD $\mathbf{X}^\top\mathbf{Y} = \mathbf{U}_{\text{svd}}\boldsymbol\Sigma\mathbf{V}_{\text{svd}}^\top$ and define the orthogonal matrix $\mathbf{Z} = \mathbf{V}_{\text{svd}}^\top\mathbf{R}^\top\mathbf{U}_{\text{svd}}$ . Then

\operatorname{tr}(\mathbf{R}^\top\mathbf{U}_{\text{svd}}\boldsymbol\Sigma\mathbf{V}_{\text{svd}}^\top) = \operatorname{tr}(\mathbf{Z}\boldsymbol\Sigma) = \sum_i Z_{ii}\,\sigma_i \le \sum_i \sigma_i ,

because every entry of an orthogonal matrix has absolute value at most 1 and the $\sigma_i$ are nonnegative. Equality holds when $\mathbf{Z} = \mathbf{I}$ , that is $\mathbf{R} = \mathbf{U}_{\text{svd}}\mathbf{V}_{\text{svd}}^\top$ . ∎

Interview

Orthogonal Procrustes

Given matched rows $\mathbf{X}, \mathbf{Y} \in \mathbb{R}^{n\times d}$ , find the rotation $\mathbf{R}$ minimizing $\lVert \mathbf{X}\mathbf{R} - \mathbf{Y}\rVert_F$ , using the result just derived. Implement procrustes(X, Y). This shows up when aligning embedding spaces across languages or model versions, comparing the representations of two networks, and in point-cloud registration.

import numpy as np

def procrustes(X, Y):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Stop 4. Embeddings that steer generation

When you type a prompt into an image generator, your words become embeddings that steer a generative model. The details differ by system, and they're worth knowing because they show how flexible the idea is.

DALL·E (Ramesh et al., 2021) turned each $256\times256$ image into a $32\times32$ grid of 1,024 discrete tokens from a learned codebook of 8,192 entries (Stop 5), then trained one large Transformer to continue a sequence of text tokens with image tokens. Image generation became next-token prediction over image tokens.
DALL·E 2 (Ramesh et al., 2022) works in CLIP's joint space (§1.8). Your prompt becomes a CLIP text embedding. A "prior" model translates it (together with the caption) into a plausible CLIP image embedding, and a diffusion decoder paints an image from that embedding.
Imagen (Saharia et al., 2022) conditioned a diffusion model on the token embeddings of a large frozen text-only language model (T5-XXL), and found that scaling the text encoder improved results more than scaling the image model, which is why the choice of text encoder matters so much.
Stable Diffusion, the open release built on the latent diffusion work of Rombach et al. (2022), does the expensive work in a compressed space. A VAE turns a $512\times512\times3$ image (786,432 numbers) into a $64\times64\times4$ grid of latent vectors (16,384 numbers, 48 times fewer), a spatial embedding. The denoiser runs on that grid. In version 1, the 77 token embeddings (each 768-dimensional) from a frozen CLIP text encoder enter every layer through cross-attention, the attention variant of §4.6. Later versions changed text encoders, with SDXL adding a second CLIP encoder and Stable Diffusion 3 adding T5 for better prompt following.

So your prompt is embedded, the image is generated as an embedding (or a grid of them), and pixels only appear at the very end. The Welch Labs guest video on 3Blue1Brown, linked with this section, is a gorgeous visual preview of how diffusion models do it.

Stop 5. Codebooks, the embedding tables of images and sound

Language is already discrete, so a token table works directly. Images and audio are continuous. Vector quantization (VQ-VAE, van den Oord et al., 2017) makes them discrete by learning a codebook, a table $\mathbf{C}_{\text{vq}}\in\mathbb{R}^{K\times d}$ of $K$ vectors (the subscript keeps it apart from the co-occurrence matrix of §1.4), and replacing each encoder output $\mathbf{z} \in \mathbb{R}^d$ with its nearest codebook entry,

k^\star = \arg\min_{k} \lVert \mathbf{z} - (\mathbf{C}_{\text{vq}})_{k,:}\rVert, \qquad \mathbf{z}_q = (\mathbf{C}_{\text{vq}})_{k^\star,:} .

In plain words, each patch of an image or slice of audio is snapped to the closest of $K$ prototype vectors, and only the index $k^\star$ is kept. Look at what that is. A codebook is an embedding table (§1.2), and choosing an entry is a nearest-neighbor search, which you can vectorize with the distance expansion of Exercise 1.1d. Images become sequences of "visual words", and audio becomes "audio words". That's how voice models and some image and video generators can use exactly the same next-token machinery as language models. The $\arg\min$ has no gradient, so training the encoder and the codebook takes a few tricks of its own, which the VQ-VAE material later in the book develops.

One codebook can only be so precise. Residual vector quantization stacks several. Quantize $\mathbf{z}$ with the first codebook, quantize what's left over (the residual) with a second, and so on, adding the chosen vectors together. Each stage refines the previous one. Neural audio codecs built this way, such as SoundStream and EnCodec, compress speech and music into a few codes per frame. EnCodec at 24 kHz produces 75 frames per second, and with 8 codebooks of 1,024 entries (10 bits each), that is $75 \times 8 \times 10 = 6{,}000$ bits per second, the bitrate of a phone call, for audio a codec language model can then generate token by token.

Stop 6. Predicting in embedding space

Here's a question that sits at the frontier. If you want a model to understand the world by predicting it, what should it predict? Predicting every pixel of the next video frame forces the model to waste effort on unpredictable details (the exact flicker of leaves). Yann LeCun's proposal, the Joint-Embedding Predictive Architecture (JEPA), is to predict the embedding of the missing or future part from the embedding of the visible part. I-JEPA does it for images (Assran et al., 2023), V-JEPA for video (Bardes et al., 2024), and world models that learn dynamics in latent space use the same idea to imagine the consequences of actions before taking them.

There's a trap, and it's worth seeing now because it shapes all of self-supervised learning. If the loss only says "make these two embeddings match", the encoder can cheat by mapping everything to the same vector. The loss is zero, and nothing is learned. This is representation collapse. Contrastive methods (§1.5, §1.8) prevent it with negatives that must be pushed apart. JEPA-style methods prevent it with tricks like a slowly updated target encoder, stop-gradients, or explicit penalties that keep each dimension's variance up, as in VICReg (Bardes, Ponce & LeCun, 2022). Watch collapse happen with a linear encoder trained by plain gradient descent, and watch a variance penalty stop it.

Python · cell 2

The matching loss alone drives the encoder toward the zero map, the trivial solution. The variance hinge only activates when a dimension's standard deviation falls below 1, and it pushes back just enough to keep every dimension alive.

The full VICReg loss adds a third term, a covariance penalty that decorrelates the dimensions, because an encoder can also collapse more subtly by making every dimension carry the same information. Non-contrastive self-supervised learning builds it, along with the momentum encoders and stop-gradients that do the same job without an explicit penalty.

Stop 7. Memory

Finally, embeddings are how AI systems remember. A retrieval-augmented system (§7.6) embeds a library of documents and fetches the relevant ones into a language model's context. Agents that run for hours or days store their experiences, notes, and tool results as embedded memories and retrieve them by similarity when they become relevant again, which is the basis of long-term memory for agents. The cosine nearest-neighbor search from §1.3 is, quite literally, how these systems recall.

Are these all really "embeddings"?

Yes. A useful definition covers every case you've seen.

An embedding is a learned mapping of things into a vector space (or a grid of vectors) such that the space's geometry, its distances, directions, and dot products, reflects what a training objective cares about.

What varies is only a handful of choices, and you can now ask them about any system you meet.

What goes in? An ID, a token, a patch, a frame, a whole document.
How is the vector produced? A lookup, a projection, a deep encoder, or a quantizer.
What shape is it? One vector per thing (retrieval, CLIP), one per token or patch (Transformers), or a spatial grid (diffusion latents).
What objective shaped it? Co-occurrence, classification, next-token prediction, contrastive pairing, reconstruction, or prediction in embedding space.
How is it used? Compared (search), decoded (generation), mapped into another space (bridges), or refined further (deeper layers).

The road ahead

question you might now have	where it's answered
What exactly goes into the token table?	Chapter 2, Tokenization
How do vectors read their context?	Chapter 3 and Chapter 4, contextual meaning and attention
How is a full language model built and trained?	Chapters 5, 6, and 8 to 13
How are sentence embedders and retrieval systems trained?	Chapter 7
How do chat models get aligned after pretraining?	Chapters 14 to 19
How do images become embeddings, and images and text share a space?	the vision part of the book, from Vision Transformers to vision-language models
How do latent spaces, codebooks, and diffusion generate images and video?	the generative part, from VAEs and discrete latents to image and video generation
How does speech become vectors and back?	the audio part, from audio encoders to voice models
What are JEPA and world models?	JEPA and world models
How do agents remember and improve themselves?	agents, memory, and self-improvement loops
What's actually inside these embedding spaces?	interpretability

🎯 Interview questions on this section

How does a language model turn its last hidden state into a next-token distribution, and what does tying change?

It computes logits $\mathbf{E}_{\text{out}}\mathbf{h}$ , one dot product per vocabulary token, then a softmax, optionally with a temperature. With tying, $\mathbf{E}_{\text{out}}$ is the input table, so the model predicts the token whose embedding best aligns with its final state, saving $Vd$ parameters and coupling input and output meanings. The same head applied to intermediate layers (the logit lens) shows how the prediction forms.

How does a vision-language model feed an image to an LLM? Give a token count.

An image encoder (often a CLIP ViT) produces patch embeddings, a projector (linear or a small MLP) maps them into the LLM's token embedding space, and they're inserted into the sequence like extra tokens. In LLaVA-1.5 a $336\times336$ image gives $24\times24 = 576$ tokens projected from width 1,024 to 4,096. Alternatives are cross-attention from text to image features (Flamingo) or a querying module that compresses the image to a fixed number of vectors, such as BLIP-2's 32 queries.

Derive the solution to orthogonal Procrustes.

Minimizing $\lVert\mathbf{XR}-\mathbf{Y}\rVert_F^2$ over orthogonal $\mathbf R$ is the same as maximizing $\operatorname{tr}(\mathbf R^\top\mathbf X^\top\mathbf Y)$ , since the other terms don't depend on $\mathbf R$ . With $\mathbf X^\top\mathbf Y = \mathbf U\boldsymbol\Sigma\mathbf V^\top$ the trace becomes $\operatorname{tr}(\mathbf Z\boldsymbol\Sigma) \le \sum\sigma_i$ for the orthogonal $\mathbf Z = \mathbf V^\top\mathbf R^\top\mathbf U$ , with equality at $\mathbf Z = \mathbf I$ , so $\mathbf R = \mathbf U\mathbf V^\top$ . It is used to align word spaces across languages and to compare representations of different models.

How does text control a diffusion model?

The prompt is encoded by a frozen text encoder (CLIP's text tower or T5) into token embeddings, which enter the denoiser through cross-attention at every layer and every denoising step (Stable Diffusion, Imagen), or through joint attention over text and image tokens (Stable Diffusion 3). DALL·E 2 instead maps a CLIP text embedding to a CLIP image embedding with a prior, then decodes it with diffusion. Imagen found that a bigger text encoder helped more than a bigger image model.

What is a VQ codebook in embedding terms, and what does residual quantization buy?

A codebook is a learned embedding table of $K$ prototype vectors. Each encoder output is replaced by its nearest entry and stored as that entry's index, so images and audio become discrete tokens that the same next-token machinery as text can model. Encoding is a nearest-neighbor search, and decoding is a table lookup. Residual quantization quantizes the leftover error with a second codebook, then a third, and adds the chosen vectors, so precision grows with the number of stages while each codebook stays small. With 8 stages of 1,024 entries at 75 frames per second, EnCodec spends $75\times8\times10 = 6{,}000$ bits per second.

What is representation collapse, and how do you prevent it?

An encoder maps every input to the same vector (complete collapse), or to a low-dimensional subspace (dimensional collapse), which trivially satisfies a "make views match" loss. You prevent it with negatives (contrastive losses such as InfoNCE), asymmetric architectures with stop-gradients or momentum target encoders (BYOL, DINO, JEPA), or explicit regularizers such as VICReg's variance hinge (keeps each dimension's standard deviation above a threshold) and covariance penalty (decorrelates dimensions).

Chapter 1 wrap-up

You started with "how do you hand the word cat to a computer?" Now you can do all of the following.

Turn text into tokens and IDs, explain why IDs, one-hot vectors, and fixed codes fail, and say what an embedding table really is, forward and backward, including bags and tied tables.
Reason fluently about dot products, cosines, distances, and high-dimensional geometry, including Johnson–Lindenstrauss and random-hyperplane hashing.
Build embeddings two ways (PPMI with SVD, and skip-gram with negative sampling), derive every gradient, and explain why they're secretly the same.
Evaluate and visualize embeddings without being fooled.
Embed users, items, categories, and graphs, and explain how encoders embed images, text, audio, and video, and how CLIP puts images and text in one space.
Name exactly where static embeddings break, measure the damage, and say why contextual models fix it.
Place every major modern system on one map of embeddings.

If you solved the exercises, you've written from scratch the core of word2vec (with an alias sampler), a recommender (two ways), an LSH index, a randomized SVD, DeepWalk, a ViT patchifier, masked mean pooling, zero-shot classification, a bias test with its permutation p-value, an optimal rotation between two spaces, and the last step of a language model.

One assumption ran quietly under all of it. We treated "cat" as one token with one row. Real language models do not read words. In GPT-2's vocabulary, "cat" at the start of a text and " cat" after a space are two different tokens with two different rows, "Cat" is a third, and a rarer word such as "carburetor" may have no row of its own and be spelled from several pieces, so its meaning has to be assembled by the layers above. What the rows of the very first table should be is a design decision, and it shapes what a model costs, how fairly it treats languages other than English, and whether it can count the letters in a word. Chapter 2 makes that decision from scratch.

🧭 What to remember

Every modern model makes, refines, compares, and decodes embeddings.
Inputs enter through tables and projections, plus embeddings of position (built in §5.4) and, for diffusion models, of the noise level.
A language model refines token embeddings layer by layer and predicts the next token by a dot product with the (often tied) token table, the tied bigram model of §1.2 with layers in between.
Bridges (often linear, and best as rotations) connect spaces, as in vision-language models and cross-lingual alignment, and Procrustes gives the optimal rotation in closed form.
Text embeddings steer generation, as in DALL·E, DALL·E 2, Imagen, and Stable Diffusion with its latent grids. Codebooks are embedding tables for images and audio, and residual quantization makes them precise.
JEPA predicts in embedding space and must avoid collapse, with negatives, momentum encoders, or penalties like VICReg's. Embeddings are also how systems retrieve and remember.
Next comes Chapter 2, which decides what the rows of the very first table are.