Embedding Layers

Chapter 1: Token Embeddings in Detail

The embedding table is deceptively simple — it's just a big matrix. But it's also one of the largest parameter blocks in the entire model. For Llama 2 7B, the embedding matrix has 131 million parameters. That's more parameters than all of GPT-2 Small (124M).

Let's understand why it's so big, how it trains, and what can go wrong.

The Size of the Table

The token embedding matrix E has shape (V, d), where V is the vocabulary size (number of unique tokens) and d is the embedding dimension (length of each vector). Every entry is a learnable floating-point number.

Parameter count = V × d. Memory = V × d × bytes_per_param.

Let's compute this for real models:

Look at Llama 3 vs Llama 2. Same embedding dimension, but Llama 3 quadrupled the vocabulary (from 32K to 128K tokens). That quadrupled the embedding table from 250 MB to 1 GB. Vocabulary size is the dominant cost driver.

Why would anyone want a bigger vocabulary? Because larger vocabularies encode text more efficiently — fewer tokens per sentence, shorter sequences, lower inference cost. But the embedding table gets proportionally larger. This is a fundamental tradeoff in model design.

How Training Works

At initialization, every row of the embedding matrix is filled with small random numbers (typically drawn from a normal distribution with std = 0.02). At this point, "cat" and "dog" have random, unrelated vectors. The model can't tell them apart any more than it can tell them from "parliament."

During training, each forward pass selects a subset of rows — one for each token in the batch. The loss gradient flows backward through the model and eventually reaches the embedding layer. But here's the key: only the rows that were selected get a gradient update. If "cat" (ID 1) appeared in the batch, row 1 gets updated. Row 5291 ("dog") is untouched.

This is fundamentally different from a linear layer, where every weight participates in every forward pass. In an embedding layer, each row trains independently — only when its token appears in a batch.

The Rare Token Problem

This selective updating creates a problem. Tokens that appear frequently — "the," "is," "a" — get updated thousands of times per epoch. Their embeddings are well-trained, nuanced, and stable. They live in exactly the right part of the vector space.

But rare tokens — "pneumonoultramicroscopicsilicovolcanoconiosis," someone's unusual name, a niche technical term — might appear 5 times in the entire training set. Five gradient updates. Their embeddings stay close to their random initialization, carrying almost no learned information.

Gradient Flow Through a Lookup

How does gradient descent work on a lookup table? Think of it this way: mathematically, looking up row i of matrix E is equivalent to computing the product e_i^T · E, where e_i is a one-hot vector (all zeros except position i). No implementation actually materializes the one-hot vector — that would waste memory — but the gradient math works out the same way.

The gradient of the loss with respect to E[i] is simply the gradient that flowed back to this layer. For rows that weren't selected, the gradient is zero (they weren't used, so changing them can't change the loss). This is why optimizers like Adam keep per-parameter statistics — rare tokens get very different update patterns from frequent ones.

python
# Demonstrating gradient flow through embedding lookup
import torch
import torch.nn as nn

V, d = 6, 3
emb = nn.Embedding(V, d)

# Forward: look up tokens 1 ("cat") and 2 ("dog")
ids = torch.tensor([1, 2])
vecs = emb(ids)           # shape: (2, 3)

# Fake loss: sum all embedding values
loss = vecs.sum()
loss.backward()

# Only rows 1 and 2 got gradients
print(emb.weight.grad[0])  # tensor([0., 0., 0.])  ← "the" untouched
print(emb.weight.grad[1])  # tensor([1., 1., 1.])  ← "cat" got gradient
print(emb.weight.grad[2])  # tensor([1., 1., 1.])  ← "dog" got gradient
print(emb.weight.grad[3])  # tensor([0., 0., 0.])  ← "sat" untouched

Chapter 2: Combining Embeddings — The Input Recipe

A single token needs more than just its word meaning. It also needs to know where it is in the sequence (position 0? position 47?) and sometimes which segment it belongs to (sentence A? sentence B?). A transformer doesn't process tokens one at a time like an RNN — it sees the whole sequence at once. Without explicit position information, it can't tell "dog bites man" from "man bites dog."

The solution: add multiple embedding vectors together. Each embedding encodes a different aspect of the input. They all live in the same d-dimensional space, and the model learns to disentangle them.

BERT's Three-Part Recipe

BERT computes the input to its transformer stack as:

Three separate lookup tables. Three separate row indices. Three vectors of the same dimension d. Added element-wise.

Total embedding parameters for BERT-Base (d=768): 30,522 × 768 + 512 × 768 + 2 × 768 = 23,440,896 + 393,216 + 1,536 = 23,835,648. The token table dominates — it's 98.3% of the embedding parameters.

GPT's Two-Part Recipe

GPT-2 and GPT-3 use a simpler recipe — no segment embedding, because they don't do sentence-pair tasks:

Modern models like Llama and GPT-4 go even further. They use Rotary Position Embeddings (RoPE) instead of a learned position table. RoPE injects position information inside the attention mechanism, not at the input. This means the only embedding table is E_token — the position information is handled differently (covered in the Positional Encoding lesson).

Hand Calculation: Adding Embeddings

Let's trace BERT's recipe by hand with d=4.

Step 1: Look up each embedding.

• E_token[2] = [0.31, -0.72, 0.15, 0.88]
• E_position[3] = [0.05, 0.12, -0.33, 0.01]
• E_segment[0] = [-0.02, 0.08, 0.11, -0.05]

Step 2: Add element-wise.

• dim 0: 0.31 + 0.05 + (-0.02) = 0.34
• dim 1: -0.72 + 0.12 + 0.08 = -0.52
• dim 2: 0.15 + (-0.33) + 0.11 = -0.07
• dim 3: 0.88 + 0.01 + (-0.05) = 0.84

This single vector now encodes three things: the word is "cat," it's at position 3, and it belongs to sentence A. All compressed into 4 numbers. The transformer layers that follow will learn to disentangle these signals.

Drag the position slider and watch. The token embedding (orange bars) stays fixed — "cat" is "cat" regardless of where it appears. The position embedding (teal bars) changes — position 0 has a different pattern than position 7. The sum below shifts accordingly.

Toggle off the position component. Now the sum looks almost identical to just the token embedding. Toggle off the token component and leave only position. Now the sum looks nothing like "cat" — it's pure positional information. The model learns to use different dimensions for different types of information, making addition work despite all three signals occupying the same vector space.

Why Add Instead of Concatenate?

The obvious alternative to addition is concatenation. Instead of adding three d-dimensional vectors, concatenate them into one 3d-dimensional vector. This keeps the information perfectly separated — no interference between token meaning and position.

So why doesn't anyone do this?

There's a deeper reason too. Addition is information-lossy in theory but sufficient in practice. With d=768 or d=4096, there are enough dimensions for the model to allocate different "channels" to different signal types. Research has shown that trained position embeddings and token embeddings are nearly orthogonal — they naturally learn to use different directions in the high-dimensional space, minimizing interference.

From Scratch: BERT's Embedding Layer

python
import torch
import torch.nn as nn

class BERTEmbedding(nn.Module):
    def __init__(self, vocab_size, max_len, d_model, n_segments=2):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb   = nn.Embedding(max_len, d_model)
        self.seg_emb   = nn.Embedding(n_segments, d_model)
        self.norm       = nn.LayerNorm(d_model)  # BERT normalizes after adding
        self.drop       = nn.Dropout(0.1)

    def forward(self, token_ids, segment_ids):
        # token_ids: (batch, seq_len)  — integer token IDs
        # segment_ids: (batch, seq_len) — 0 or 1
        seq_len = token_ids.size(1)
        positions = torch.arange(seq_len, device=token_ids.device)

        # Three lookups, one addition
        x = self.token_emb(token_ids)    # (batch, seq, d)
        x = x + self.pos_emb(positions)  # broadcast: (seq, d) → (batch, seq, d)
        x = x + self.seg_emb(segment_ids)
        return self.drop(self.norm(x))

# Usage
emb = BERTEmbedding(vocab_size=30522, max_len=512, d_model=768)
tokens = torch.randint(0, 30522, (2, 128))  # batch=2, seq=128
segs   = torch.zeros(2, 128, dtype=torch.long)
out    = emb(tokens, segs)  # (2, 128, 768)

python
# HuggingFace: it's already inside the model
from transformers import BertModel
model = BertModel.from_pretrained("bert-base-uncased")

# The three embedding tables live here:
print(model.embeddings.word_embeddings.weight.shape)      # (30522, 768)
print(model.embeddings.position_embeddings.weight.shape)  # (512, 768)
print(model.embeddings.token_type_embeddings.weight.shape)# (2, 768)

Chapter 3: Patch Embeddings — Images as Token Sequences

Transformers process sequences. Text is a sequence of tokens — that's natural. But images aren't sequences. They're 2D grids of pixels. How do you feed an image into a transformer?

You could treat every pixel as a token. A 224×224 RGB image has 150,528 pixel values. But attention is O(n²) in sequence length — computing attention over 150K tokens is computationally impossible. Even with modern hardware, that's 22.6 billion attention scores per layer.

The Vision Transformer (ViT) solved this with a simple trick: chop the image into a grid of non-overlapping patches, treat each patch as a "word," and run the same transformer architecture. A 224×224 image with 16×16 patches gives 196 "tokens" — a manageable sequence that attention can process.

The Mechanics

Here's exactly what happens, step by step:

Notice the coincidence: for ViT-Base with 16×16 patches, each flattened patch has 16 × 16 × 3 = 768 values — exactly the embedding dimension d=768. So the linear projection is a square matrix (768×768). This wasn't a design accident.

Hand Calculation: A Tiny Image

Let's work through a tiny example. Grayscale image (1 channel), 6×6 pixels, patch size 3×3.

Here's our 6×6 image (pixel values 0-9):

Patch 0 (top-left 3×3 block):

• Pixels: [[1, 2, 3], [4, 5, 0], [7, 8, 1]]
• Flattened: [1, 2, 3, 4, 5, 0, 7, 8, 1] — a 9-dimensional vector

Patch 1 (top-right):

• Pixels: [[7, 8, 9], [6, 5, 4], [3, 2, 1]]
• Flattened: [7, 8, 9, 6, 5, 4, 3, 2, 1]

Patch 2 (bottom-left):

• Flattened: [2, 3, 4, 5, 6, 7, 8, 9, 0]

Patch 3 (bottom-right):

• Flattened: [8, 7, 6, 5, 4, 3, 2, 1, 0]

Linear projection: Weight matrix W is (9, 4). Let's project Patch 0:

For simplicity, suppose the projection of [1, 2, 3, 4, 5, 0, 7, 8, 1] through W gives us [0.42, -0.18, 0.73, 0.05]. That 4D vector is Patch 0's patch embedding. It's now a "token" that attention can process, just like the word "cat" in a language model.

After projecting all 4 patches, we have a sequence of length 4, each element a d=4 vector. The transformer processes this exactly like a 4-token text sequence.

The Conv2d Trick

The flatten-then-project operation can be implemented as a single convolution with kernel_size = patch_size and stride = patch_size. This is not just clever engineering — it's mathematically identical.

A convolution with kernel_size=16 and stride=16 slides a 16×16 window across the image, moving exactly 16 pixels each step (no overlap). At each position, it computes a weighted sum of the 768 pixel values (16×16×3) to produce one output value per filter. With d_model filters, each position produces a d-dimensional vector. Each position corresponds to one patch.

python
# Patch embedding from scratch: reshape + linear
import torch
import torch.nn as nn

class PatchEmbedScratch(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, d_model=768):
        super().__init__()
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2  # 196
        patch_dim = in_channels * patch_size * patch_size  # 768
        self.proj = nn.Linear(patch_dim, d_model)

    def forward(self, x):
        # x: (batch, channels, H, W)
        B, C, H, W = x.shape
        p = self.patch_size
        # Reshape: (B, C, H, W) → (B, n_patches, patch_dim)
        x = x.unfold(2, p, p).unfold(3, p, p)  # (B, C, nH, nW, p, p)
        x = x.contiguous().view(B, C, -1, p * p)  # (B, C, n_patches, p²)
        x = x.permute(0, 2, 1, 3).reshape(B, -1, C * p * p)  # (B, n_patches, patch_dim)
        return self.proj(x)  # (B, n_patches, d_model)

python
# The Conv2d trick: mathematically identical, faster on GPU
class PatchEmbedConv(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, d_model=768):
        super().__init__()
        # One conv: kernel=16, stride=16 → no overlap, each position = one patch
        self.proj = nn.Conv2d(in_channels, d_model,
                              kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # x: (B, 3, 224, 224)
        x = self.proj(x)  # (B, d_model, 14, 14)
        x = x.flatten(2)   # (B, d_model, 196)
        x = x.transpose(1, 2)  # (B, 196, d_model) — sequence of patch embeddings
        return x

# Both produce identical shapes:
img = torch.randn(1, 3, 224, 224)
scratch = PatchEmbedScratch()
conv = PatchEmbedConv()
print(scratch(img).shape)  # (1, 196, 768)
print(conv(img).shape)     # (1, 196, 768)

Patch Size: The Resolution-Speed Tradeoff

Patch size is the single most impactful hyperparameter in ViT. It controls both the resolution of the representation and the computational cost.

Going from patch size 16 to 8 quadruples the number of patches, which increases the attention cost by 16× (since attention is quadratic). This is why ViT-B/16 (patch size 16) is far more common than ViT-B/8 in practice — the accuracy gain from smaller patches rarely justifies the 16× compute increase.

Click different patches and watch the flattened vector change. Patches showing sky have low-contrast, similar pixel values — their vectors are "boring." Patches with edges or objects have high-variance pixels — their vectors are more "interesting." The linear projection learns to extract the patterns that matter for the downstream task.

The [CLS] Token

ViT adds one extra token at the beginning of the sequence: the [CLS] token. This is a learnable embedding (not from any patch) that aggregates information from the entire image through attention. After the transformer, the [CLS] token's final hidden state is used for classification.

So the final input sequence is actually 197 tokens for a 224×224 image with 16×16 patches: 1 [CLS] + 196 patch embeddings. Each token also gets a learnable position embedding added (similar to BERT's recipe from Chapter 2), so the model knows which patch is top-left vs. bottom-right.

python
# ViT's full embedding: patch projection + [CLS] + position
class ViTEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, d_model=768):
        super().__init__()
        n_patches = (img_size // patch_size) ** 2  # 196

        # Patch embedding via Conv2d
        self.patch_emb = nn.Conv2d(3, d_model,
                                   kernel_size=patch_size, stride=patch_size)

        # Learnable [CLS] token — shape (1, 1, d)
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Position embeddings for CLS + all patches
        self.pos_emb = nn.Parameter(torch.randn(1, n_patches + 1, d_model) * 0.02)

    def forward(self, x):
        B = x.shape[0]
        # Patch embeddings: (B, 3, 224, 224) → (B, 196, 768)
        patches = self.patch_emb(x).flatten(2).transpose(1, 2)

        # Prepend [CLS] token
        cls = self.cls_token.expand(B, -1, -1)  # (B, 1, 768)
        x = torch.cat([cls, patches], dim=1)     # (B, 197, 768)

        # Add position embeddings
        x = x + self.pos_emb  # (B, 197, 768)
        return x

Chapter 4: Tied Embeddings

The input embedding maps token IDs to vectors. The output layer maps vectors back to token scores. Both are matrices of shape vocabulary × dimension. What if they were the same matrix? Using it forward to embed, using its transpose to predict?

That's weight tying, and it cuts your embedding parameters in half while often making the model better.

The Output Layer Problem

A language model's final job is to predict the next token. It has a hidden state h — a vector of dimension d that summarizes everything the model has read so far. It needs to turn h into a score for every token in the vocabulary. Token 0 gets a score, token 1 gets a score, all the way up to token V-1.

The standard approach: multiply h by a weight matrix W_out of shape (V, d). The result is a vector of V logits — one score per vocabulary token. The highest logit wins.

But wait. W_out has shape (V, d). The input embedding matrix E also has shape (V, d). Both map between a vocabulary-sized space and a d-dimensional space. They're doing mirror-image jobs.

How Tying Works

Weight tying sets W_out = E. The same matrix serves double duty:

• Forward pass (embedding): look up row i of E to get token i's embedding vector.
• Output pass (prediction): compute logits = h · E^T. The dot product of h with each row of E gives the score for that token.

The dot product h · E[i] measures how similar the hidden state is to token i's embedding. High similarity = high score = model predicts that token. This creates a beautiful symmetry: tokens whose embeddings are close in the embedding space produce similar predictions at the output.

Hand Calculation: Tied Output Projection

Let's trace through a concrete example. V = 4 tokens, d = 3 dimensions.

Embedding matrix E (4×3):

The model has processed the sentence and produced a hidden state h = [0.6, 0.1, 0.3]. We now compute logits = h · E^T, which means taking the dot product of h with each row of E:

Token 0 ("the"):

• 0.6 × 0.5 + 0.1 × 0.3 + 0.3 × (-0.1)
• = 0.30 + 0.03 - 0.03 = 0.30

Token 1 ("cat"):

• 0.6 × 0.8 + 0.1 × (-0.2) + 0.3 × 0.4
• = 0.48 - 0.02 + 0.12 = 0.58

Token 2 ("sat"):

• 0.6 × 0.1 + 0.1 × 0.9 + 0.3 × 0.3
• = 0.06 + 0.09 + 0.09 = 0.24

Token 3 ("on"):

• 0.6 × (-0.3) + 0.1 × 0.5 + 0.3 × 0.6
• = -0.18 + 0.05 + 0.18 = 0.05

Logits: [0.30, 0.58, 0.24, 0.05]. Token 1 ("cat") has the highest score. The hidden state h is most similar to the embedding for "cat," so the model predicts "cat" as the next token.

The Parameter Savings

Let's count parameters for a real model. Llama 2 7B: V = 32,000 tokens, d = 4,096 dimensions.

That's 131 million fewer parameters — about 250 MB of memory saved in FP16 precision. For larger vocabularies (Llama 3's 128K tokens), the savings are even more dramatic: 128,000 × 4,096 = 524M params saved.

But tying isn't only about saving memory. The shared representation acts as a constraint.

Who Ties and Who Doesn't

Weight tying is common but not universal:

The trend: smaller models benefit more from tying (the embedding parameters are a bigger fraction of total). Larger models can afford separate matrices and sometimes achieve marginally better performance without tying.

The Simulation

Below, we visualize the relationship between input embeddings and output projections. Toggle "tied" to see the matrices link together. The hidden state projects through the embedding matrix to produce scores — tokens closer to h in embedding space get higher logits.

In Code

python
import torch
import torch.nn as nn

class TiedLM(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        # Output projection shares the embedding weight
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.embed.weight  # THE TIE

    def forward(self, token_ids, hidden):
        # Input: embed tokens
        x = self.embed(token_ids)       # (B, T) -> (B, T, d)
        # ... transformer layers process x into hidden ...
        # Output: logits = hidden @ E.T
        logits = self.lm_head(hidden)  # (B, T, d) -> (B, T, V)
        return logits

# Verify they share memory
model = TiedLM(32000, 4096)
print(model.embed.weight.data_ptr() == model.lm_head.weight.data_ptr())
# True — same tensor in memory

The critical line is self.lm_head.weight = self.embed.weight. After this assignment, both layers reference the exact same tensor in memory. Gradient updates from the output loss flow into the same parameters that define the input embeddings. One gradient update improves both embedding and prediction simultaneously.

Chapter 5: Embedding Scaling

You've combined token and position embeddings by addition. But what if one is much larger than the other? If position embeddings have values around 1.0 and token embeddings have values around 0.02, the position signal dominates. The model knows where every token is but barely knows what it is.

This is exactly what happens with the original Transformer's initialization. And the fix is a single multiplication.

The Scale Mismatch

When we initialize an embedding matrix, each entry is drawn from a distribution with standard deviation roughly 1/√d. For a model with d = 512, that means each element is about ±0.044.

How big is an entire embedding vector? Each of its d elements has variance 1/d, so the vector's squared magnitude (sum of squared elements) is approximately d × (1/d) = 1. The magnitude of a token embedding is roughly √1 = 1.0.

Now consider sinusoidal position embeddings. Each element is sin(…) or cos(…), so values range from -1 to +1. The squared magnitude is approximately d × 0.5 = d/2. The magnitude is roughly √(d/2).

For d = 512: position magnitude ≈ √256 = 16.0. Token magnitude ≈ 1.0. When we add them, position contributes 16× more than token identity. The model is almost entirely position, with a tiny whisper of "what token is this?"

The Fix: Scale by √d

After multiplying by √d, each element goes from std ≈ 1/√d to std ≈ 1. The vector's magnitude goes from ≈1 to ≈ √d. Now it matches the positional embedding magnitude.

Hand Calculation: Magnitude Comparison

Let d = 512. A token embedding vector e has d elements, each with std = 1/√512 ≈ 0.0442.

Without scaling:

• Expected element magnitude: 0.0442
• Expected vector magnitude: √(d × (1/√d)²) = √(512 × 1/512) = √1 = 1.0

Sinusoidal position embedding:

• Element values: sin/cos in [-1, +1], expected squared value ≈ 0.5
• Expected vector magnitude: √(d × 0.5) = √256 = 16.0

Ratio without scaling: position/token = 16.0/1.0 = 16:1. Position dominates.

With scaling (multiply e by √512 ≈ 22.6):

• Each element: 0.0442 × 22.6 ≈ 1.0
• Vector magnitude: 1.0 × 22.6 = 22.6
• Ratio: 16.0/22.6 ≈ 0.7:1. Balanced!

Now both signals contribute comparably. The model can distinguish both what a token is and where it is from the very first layer.

A Concrete Example

Let's use d = 4 for a tiny example we can trace completely.

Token embedding (initialized with std = 1/√4 = 0.5):

Magnitude: √(0.09 + 0.25 + 0.04 + 0.16) = √0.54 ≈ 0.735

Positional embedding (sinusoidal):

Magnitude: √(0.71 + 0.29 + 0.83 + 0.18) = √2.01 ≈ 1.418

Without scaling: combined = e + p

The result is dominated by the position values (0.84, 0.54, 0.91, -0.42). The token's contribution is barely visible — the 0.3 gets swamped by 0.84, the -0.5 nearly cancels with 0.54.

With scaling: e_scaled = e × √4 = e × 2

combined = e_scaled + p

Now both signals are clearly visible. The -1.0 from the token is preserved alongside the 0.54 from position, giving -0.46 instead of a washed-out 0.04.

When Scaling Is and Isn't Used

The pattern: scaling is mainly needed when fixed (sinusoidal) position embeddings are added to small-initialized token embeddings. Modern LLMs using RoPE don't add position to the embedding at all — they rotate the query and key vectors. No addition means no magnitude mismatch.

The Simulation

Adjust the scaling factor below to see how token and position signals combine. At low scale, position dominates (all tokens look similar in the combined representation). At √d, they balance. At very high scale, token dominates and position information is lost.

In Code

python
import torch
import torch.nn as nn
import math

class ScaledEmbedding(nn.Module):
    """Original Transformer embedding with sqrt(d) scaling."""
    def __init__(self, vocab_size, d_model, max_len=5000):
        super().__init__()
        self.d_model = d_model
        self.tok_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = self._sinusoidal(max_len, d_model)

    def forward(self, token_ids):
        seq_len = token_ids.size(1)
        tok = self.tok_embed(token_ids)              # (B, T, d)
        tok = tok * math.sqrt(self.d_model)          # SCALE!
        pos = self.pos_embed[:seq_len].unsqueeze(0)  # (1, T, d)
        return tok + pos                              # balanced addition

The single line tok * math.sqrt(self.d_model) makes the difference between a model that can barely distinguish tokens and one where token identity and position are balanced from the start.

Chapter 6: The Embedding Space Explorer

Everything we've learned — token IDs mapped to vectors, embeddings trained by gradient descent, parameter sharing, magnitude balancing — converges into one central question: what does the learned embedding space actually look like?

Random initialization scatters tokens across the space with no structure. Training sculpts that space: synonyms cluster together, antonyms push apart, and semantic relationships emerge as geometric patterns. The explorer below lets you watch this process unfold in real time.

What You'll See

The canvas shows a 2D projection of a simulated embedding space. Each dot is a token. At step 0 (random initialization), the dots are scattered with no pattern. As training progresses, semantic clusters emerge. Animals drift toward animals. Colors group with colors. The structure you see is the model discovering meaning.

Play with the controls:

• Vocabulary theme: Choose which tokens populate the space. Each theme has natural clusters the training will discover.
• Embedding dimension: Higher dimensions give the space more room. Clusters tighten because there are more directions to separate them. Lower dimensions force compromises — clusters overlap.
• Training steps: Slide or animate to watch the embeddings evolve. The first few steps show dramatic reorganization. Later steps refine the structure.
• Show similarity lines: Toggle lines between tokens with high cosine similarity. After training, these connect semantically related tokens.
• Show clusters: Draw convex hulls around discovered clusters.

What the Training Is Doing

In a real language model, embedding training happens implicitly through the language modeling objective: predict the next token. Tokens that appear in similar contexts develop similar embeddings because the gradients push them in similar directions.

Our simulation mimics this with a simplified attraction/repulsion model. Tokens in the same semantic category attract each other (their embeddings move closer). Tokens in different categories repel slightly. The result approximates what real language model training produces.

Three things to notice as training progresses:

• Cluster formation (steps 0-100): Initially random tokens rapidly separate into groups. This is the "big picture" phase where the model learns major categories.
• Cluster tightening (steps 100-300): Within-cluster tokens move closer together. The model refines distinctions: "cat" and "dog" aren't just "animals" — they're specific animals.
• Fine structure (steps 300-500): Subtle sub-clusters appear. Among animals, pets cluster separately from wildlife. Among colors, warm colors separate from cool colors.

Chapter 7: Subword Embeddings

The word "transformerization" might appear 3 times in your entire training set. Its embedding gets 3 gradient updates — barely better than random. But break it into subwords: "transform," "er," "ization." Each piece appears thousands of times. The composed representation from well-trained subwords is far better than a single undertrained whole-word embedding.

This is the core insight of subword tokenization, and it's why every modern language model uses it.

The Whole-Word Problem

Suppose we give every word its own embedding. English has roughly 170,000 words in common use. Add technical terms, names, misspellings, and multilingual support, and you easily hit a million. That's a million rows in the embedding matrix, most of them barely trained.

Worse: a new word the model has never seen gets no embedding at all. The entire model breaks on a single unfamiliar word. This is the out-of-vocabulary (OOV) problem, and it plagued NLP for decades.

BPE: Byte Pair Encoding

Byte Pair Encoding (Sennrich et al., 2016) solves both problems with a simple algorithm. Start with a character-level vocabulary (every letter is a token). Then repeatedly merge the most frequent adjacent pair into a new token.

Hand Calculation: BPE Merges

Let's run BPE on a tiny corpus. Our training text:

text
"the cat sat on the mat the cat"

Step 0: Character tokens

['t','h','e',' ','c','a','t',' ','s','a','t',' ','o','n',' ','t','h','e',' ','m','a','t',' ','t','h','e',' ','c','a','t']

Count all adjacent pairs:

Merge 1: Most frequent pair is "a"+"t" (count 4). Create new token "at".

Result: ['t','h','e',' ','c','at',' ','s','at',' ','o','n',' ','t','h','e',' ','m','at',' ','t','h','e',' ','c','at']

Merge 2: Now "t"+"h" appears 3 times. Create "th".

Result: ['th','e',' ','c','at',' ','s','at',' ','o','n',' ','th','e',' ','m','at',' ','th','e',' ','c','at']

Merge 3: "th"+"e" appears 3 times. Create "the".

Result: ['the',' ','c','at',' ','s','at',' ','o','n',' ','the',' ','m','at',' ','the',' ','c','at']

After just 3 merges, "the" is a single token. High-frequency words collapse quickly. Rare words stay decomposed into common pieces.

Why Subword Embeddings Are Better

Consider the word "unhappiness" (appears rarely) versus its subwords:

The subword approach gives us three well-trained embeddings instead of one barely-trained one. The transformer's attention layers then compose these subword embeddings into a word-level representation. Attention learns that "un" means negation, "happi" carries the core meaning, and "ness" marks a noun. This compositional understanding transfers to every word with these subwords — "unhelpfulness," "happiness," "sadness" all benefit.

The Vocabulary Size Tradeoff

The number of BPE merges determines the vocabulary size. This is a critical hyperparameter with opposing forces:

Most modern LLMs settle on V = 32K-128K as the sweet spot. Llama uses 32K. GPT-4 uses ~100K. Llama 3 moved to 128K to improve multilingual coverage (more languages means more unique subwords needed).

The Simulation

Type any word below to see how it gets tokenized under different vocabulary sizes. The simulation shows which subwords are well-trained (green, high frequency) versus barely trained (red, low frequency). Notice how smaller vocabularies split words into more pieces, but each piece is individually better trained.

In Code

python
# Simplified BPE from scratch
def get_pairs(tokens):
    """Count all adjacent pairs."""
    pairs = {}
    for i in range(len(tokens) - 1):
        pair = (tokens[i], tokens[i+1])
        pairs[pair] = pairs.get(pair, 0) + 1
    return pairs

def merge(tokens, pair, new_token):
    """Replace every occurrence of pair with new_token."""
    result = []
    i = 0
    while i < len(tokens):
        if i < len(tokens)-1 and tokens[i]==pair[0] and tokens[i+1]==pair[1]:
            result.append(new_token)
            i += 2
        else:
            result.append(tokens[i])
            i += 1
    return result

# Run BPE
text = "the cat sat on the mat"
tokens = list(text)            # start with characters
num_merges = 10

for _ in range(num_merges):
    pairs = get_pairs(tokens)
    if not pairs: break
    best = max(pairs, key=pairs.get)
    new_tok = best[0] + best[1]
    tokens = merge(tokens, best, new_tok)
    print(f"Merge: '{best[0]}' + '{best[1]}' -> '{new_tok}'  Tokens: {tokens}")

python
# Production tokenizers: tiktoken (GPT) and sentencepiece (Llama)
import tiktoken
enc = tiktoken.encoding_for_model("gpt-4")
tokens = enc.encode("transformerization")
print([enc.decode([t]) for t in tokens])
# ['transform', 'er', 'ization']

print(len(tokens))  # 3 subword tokens

# Compare: a common word is a single token
print(enc.encode("the"))   # [1820] — one token
print(enc.encode("hello")) # [15339] — one token

The Embedding Quality Hierarchy

Not all embeddings are created equal. The number of gradient updates a token receives during training directly determines how well its embedding represents its meaning:

This is why the worst case for subword tokenization is still decent: even a word split into individual characters uses embeddings that have each seen millions of training examples. The model can still compose meaning from well-trained character embeddings — it just takes more layers of attention to do so.

Chapter 8: Embedding Strategy Arena

You've learned six different ways to build embedding layers: small vocabularies with tied weights, large vocabularies with untied weights, subword BPE, character-level, and two patch sizes for vision. Each has tradeoffs in parameter count, embedding quality, sequence length, and compute cost. But which one wins on a real task?

That depends entirely on the task. The arena below lets you race all six strategies head-to-head across four different challenges. No single strategy dominates everywhere — that's the whole point of understanding the tradeoffs.

The Six Strategies

The Four Tasks

Each task favors different strengths. Watch which strategies rise and fall:

• Language Modeling: Predict the next token. Large vocabularies and subword BPE shine — fewer tokens per sentence means less attention cost and better long-range dependencies. Character-level struggles with sequence length.
• Sentence Embedding: Produce a single vector summarizing a sentence. Tied weights help because the shared representation creates consistent semantics. Character-level pays a heavy attention cost for long sequences.
• Image Classification: Classify a 224×224 image. Only patch strategies apply here. Patch 16 gives more detail (196 tokens) but costs more. Patch 32 is faster (49 tokens) but misses fine features.
• Multilingual: Handle text in 100+ languages. Large vocabularies cover more scripts. Character-level handles any script but makes very long sequences. Small vocabularies fail on non-Latin scripts.

Reading the Results

No strategy wins every task. That's the deepest lesson here:

Chapter 9: Cheat Sheet

Everything from this lesson, compressed into a reference you can come back to.

Decision Flowchart

Component Catalog

PyTorch Quick Reference

python
import torch
import torch.nn as nn
import math

# 1. Basic token embedding
emb = nn.Embedding(num_embeddings=32000, embedding_dim=768)
vectors = emb(torch.tensor([42, 1337]))  # shape: (2, 768)

# 2. Weight tying
lm_head = nn.Linear(768, 32000, bias=False)
lm_head.weight = emb.weight  # same tensor in memory

# 3. Embedding scaling (original Transformer)
scaled = emb(ids) * math.sqrt(768)  # match sinusoidal magnitude

# 4. Patch embedding (ViT)
patch_emb = nn.Conv2d(3, 768, kernel_size=16, stride=16)
patches = patch_emb(img).flatten(2).transpose(1, 2)  # (B, 196, 768)

# 5. BERT-style combined embedding
x = emb(token_ids) + pos_emb(positions) + seg_emb(segment_ids)

# 6. Subword tokenization (production)
import tiktoken
enc = tiktoken.encoding_for_model("gpt-4")
tokens = enc.encode("transformerization")  # ['transform', 'er', 'ization']

Key Numbers to Remember

Connections

Embedding layers are the first thing that happens in every neural network that processes discrete inputs. Understanding them unlocks everything downstream:

• Normalization — LayerNorm is typically applied right after the embedding addition. Why? Because adding three vectors of different scales can create magnitude issues that normalization fixes.
• Optimizers — Embedding layers have unique optimization challenges: sparse gradients (only selected rows update), rare tokens (few gradient steps), and the weight tying constraint. Adam's per-parameter statistics handle these naturally.
• Transformers — The embedding output is what enters the transformer stack. Every concept in the transformer lesson — attention, feed-forward networks, residual connections — operates on the vectors that embedding layers produce.
• Vision Transformer — Patch embeddings are how ViT converts images into the sequence format that transformers require. The tradeoffs of patch size directly affect model quality and efficiency.