Token embeddings

Token embeddings, from string to (B, T, 768)

Overview

This lesson introduces token embeddings and how they turn discrete tokens into continuous vector representations.

From characters to tokens

Before a language model can look up an embedding, it has to decide what counts as a "token." Different model families make different choices, and those choices shape how the model sees text, what it groups together, what it splits, where its weak spots are.

The playground below runs three real tokenizers (GPT-2's byte-level BPE, BERT's WordPiece, and Llama's SentencePiece) entirely in your browser. Type once, and watch how each one carves up the same input.

A few things to try:

• Indented code, count how many tokens each model spends on the leading whitespace alone.
• Glitch tokens like SolidGoldMagikarp, a known artifact in GPT-2's vocabulary; see how it tokenizes.
• Non-English text (Hindi, Mandarin), most tokenizers fall back to byte-level fragments, dramatically inflating the token count.
• Numbers like 1234567890, different tokenizers split these very differently.

The colors are stable: the same token ID always gets the same color, so when the same colored block reappears, that's the model reusing one vocabulary entry.

Where we are in the build

Lesson 1 left you with config.py, a single dataclass that says "GPT-2 small is 12 layers, 12 heads, 768-dim embeddings, vocab of 50257." Today we fill in the very first thing the model touches: the path from a string of English to a tensor of floats.

This lesson adds two files to your scaffold:

1. tokenizer.py, turn a string into a list of integer token ids in [0, 50257).
2. embeddings.py, turn each integer token id into a 768-dimensional learnable vector.

After this lesson, you can call encode("Hello world") and get back [15496, 995], then feed those integers into a TokenEmbedding(config) and get back a (1, 2, 768) tensor. That tensor is the input to the very first transformer block, everything else in the model operates on these vectors.

The naive solution, one-hot vectors

Before we use a learned embedding table, let's do the dumbest possible thing and see why it can't work. Each token id $i \in [0, 50257)$ becomes a vector with a single 1 at position $i$ and zeros everywhere else:

Now imagine a batch of 32 sequences of 1024 tokens each. The input tensor is (32, 1024, 50257) of float32:

Per batch. Just to represent the input. Before we've done a single matrix multiply.

That's just the memory problem. The deeper problem is that one-hot vectors carry no information about token similarity. The vector for "cat" is exactly as far from "dog" as it is from "subpoena", they're all orthogonal unit vectors, distance $\sqrt{2}$ from each other. The model can't generalise from "the cat sat" to "the dog sat" without learning every word independently. We've thrown away every shred of structure that would help the model learn.

The motivated solution, learned dense embeddings

The fix is to replace the 50257-dimensional sparse vector with a dense vector of 768 floats, and to learn those 768 floats by gradient descent, alongside the rest of the model. Tokens that play similar roles (synonyms, conjugations, related concepts) end up with similar vectors. The model picks the geometry that helps it predict the next token.

Mechanically, an embedding layer is a single matrix $E \in \mathbb{R}^{V \times C}$ where $V = 50257$ and $C = 768$. To look up token $i$, you take row $i$ of $E$. That's it, no matrix multiply, just an indexing operation that PyTorch implements as nn.Embedding.

The parameter count works out to:

For GPT-2 small (124M total), about 31% of the entire model lives in this one table. Embeddings are not a side-show, they are the foundation, and a substantial fraction of the optimiser's job is learning them.

Tokenizer code

# nanogpt-from-scratch/tokenizer.py
 
import tiktoken
 
_enc = tiktoken.get_encoding("gpt2")
 
 
def encode(text: str) -> list[int]:
    """Map a string to a list of token ids in [0, 50257)."""
    return _enc.encode(text)
 
 
def decode(ids: list[int]) -> str:
    """Map token ids back to a string."""
    return _enc.decode(ids)
 
 
def vocab_size() -> int:
    return _enc.n_vocab   # 50257

We do not reimplement BPE, that would be a course of its own. tiktoken is OpenAI's official library and gives us the exact same vocabulary and merges as GPT-2. The point of this file is to have a single import path (from tokenizer import encode, decode) that the rest of the project uses.

Embedding code

# nanogpt-from-scratch/embeddings.py
 
import torch
import torch.nn as nn
 
from config import GPTConfig
 
 
class TokenEmbedding(nn.Module):
    """Map token ids -> dense vectors. (B, T) -> (B, T, n_embd)."""
 
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.weight = nn.Embedding(config.vocab_size, config.n_embd)
 
    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        # idx: (B, T) int64 token ids
        # returns: (B, T, n_embd) float32
        return self.weight(idx)

A few things to notice:

• No bias. Embeddings don't get added, they get looked up. There is no + b step.
• Initialisation matters. PyTorch initialises nn.Embedding with $\mathcal{N}(0, 1)$. nanoGPT and the GPT-2 paper use $\mathcal{N}(0, 0.02)$, closer to the small variance the rest of the model expects. We'll set this in Lesson 10 when we assemble the full model.
• The same matrix shows up at the output. When the model finishes producing a (B, T, 768) hidden state, it has to project back to (B, T, 50257) to score every token. GPT-2 ties the output projection to this same embedding matrix, one set of 38.6M weights, used twice. Halves the embedding parameter cost and tends to give slightly better perplexity. We'll wire this in Lesson 10.

Math by typing it

The embedding lookup is one matrix and one indexing operation:

Looked up over a batch of token ids $\text{idx} \in \mathbb{Z}^{B \times T}$:

Equivalent, and how some pedagogy presents it, is to one-hot $\text{idx}$ to a (B, T, V) tensor and matmul by $E$. The result is identical, but nn.Embedding does the indexing form so you skip allocating the giant sparse tensor.

Memory accounting for GPT-2 small:

Run it on your laptop

# smoke_test.py, at the project root, run with `python smoke_test.py`
 
import torch
 
from config import GPTConfig
from tokenizer import encode, decode, vocab_size
from embeddings import TokenEmbedding
 
# Tokenize.
ids = encode("Hello world")
print("token ids:", ids)               # [15496, 995]
print("round-trip:", decode(ids))      # "Hello world"
print("vocab size:", vocab_size())     # 50257
 
# Embed.
config = GPTConfig()
embed = TokenEmbedding(config)
idx = torch.tensor([ids])              # shape (1, 2)
out = embed(idx)
print("embedding shape:", out.shape)   # torch.Size([1, 2, 768])
print("dtype:", out.dtype)             # torch.float32
 
# Sanity: same id always yields the same vector.
assert torch.allclose(embed(torch.tensor([[15496]])),
                      embed(torch.tensor([[15496]])))
 
# Sanity: the embedding has the parameter count we expect.
n_params = sum(p.numel() for p in embed.parameters())
print(f"embedding params: {n_params:,}")   # 38,597,376

If the parameter count prints 38,597,376, you've matched GPT-2 small's embedding table exactly.

What you can now do

• Encode arbitrary English text to a sequence of integer token ids using GPT-2's exact vocabulary.
• Convert any sequence of token ids into a (B, T, n_embd) tensor of learnable vectors, ready for the rest of the model.
• Reason about the parameter cost of the embedding table, and why it's the single largest weight matrix in the model until Lesson 8 introduces the FFN.
• Recognise the weight tying trick that GPT-2, GPT-3, and most subsequent LLMs use to avoid paying for the embedding twice.

Quiz

Practice it yourself

Feel why one-hot fails and how weight tying pays off. Real Python, hidden tests, no setup.

What's next

You can turn text into vectors. But there's a problem the playground above hints at: token ids are just integers. The vector for the first "the" in a sentence is identical to the vector for the second "the", the model has no idea which token came first. Lesson 3 fixes that by extending embeddings.py with a PositionalEmbedding layer that adds an order signal to every token's vector before the first transformer block sees it.