Tokens and Embeddings: Turning Text into Vectors

Where the course begins

This is the first lesson of a course that builds a decoder-only transformer — a small GPT — one mechanism at a time, in PyTorch, until the final lesson trains a model that generates text. There is no previous lesson to recap; there is only the running example we will extend from here to the end: a character-level GPT trained on a small text corpus, the same setup as Andrej Karpathy's nanoGPT.⁶ Everything else in the course — attention, multiple heads, positional encoding, the transformer block, the training loop — assumes the machinery this lesson builds. A transformer never sees text. It sees a tensor of floating-point numbers of a fixed width the original paper calls d_model, and the first job of any language model is to turn one into the other faithfully.¹

Two terms to fix at the outset. A token is the atomic unit of text the model reads — a character, a subword, or a whole word, depending on the scheme. The vocabulary is the finite set of distinct tokens the model knows; every token is identified by an integer index into it. The entire pipeline below exists to carry a string into that integer space and then into a continuous vector space the network can optimize over.

From characters to integers

Start with the simplest tokenizer that works: one token per character. Given a corpus, the vocabulary is just the sorted set of distinct characters that appear in it, and each character maps to its position in that set. This is the scheme Andrej Karpathy's nanoGPT uses for its tiny-Shakespeare demo, where the corpus resolves to a vocabulary of 65 characters — the letters, digits, punctuation, and whitespace present in the text.⁶ Two lookup tables, string-to-integer and integer-to-string, are the whole tokenizer.

text = open("input.txt").read()            # the tiny-shakespeare corpus
chars = sorted(set(text))                  # every distinct character, in order
vocab_size = len(chars)                     # 65 for tiny shakespeare
stoi = {c: i for i, c in enumerate(chars)}  # string -> integer id
itos = {i: c for i, c in enumerate(chars)}  # integer id -> string

encode = lambda s: [stoi[c] for c in s]     # "Hi" -> [20, 47]
decode = lambda ids: "".join(itos[i] for i in ids)

encode("Hello")   # -> [20, 43, 50, 50, 53]

Character tokenization is honest and trivial to invert, but it pays for that simplicity in sequence length: every character is a step the model must attend over, so a paragraph becomes hundreds of tokens. Production language models make the opposite trade. GPT-2 uses a byte-level byte-pair encoding (BPE) tokenizer with a vocabulary of 50,257 tokens and a context window of 1,024 tokens.³ BPE, introduced for neural machine translation by Sennrich, Haddow, and Birch, adapts a data-compression algorithm to segmentation: it encodes rare and unknown words as sequences of frequent subword units, so the model gets an open vocabulary without a token for every word it might ever meet.²

The mechanics are visible in GPT-2's released tokenizer. A bytes_to_unicode map assigns each of the 256 byte values a printable character, and a ranked table of merges is applied greedily, highest-priority pair first, until no ranked pair remains; the base of 256 bytes plus the learned merges is what produces the final vocabulary.⁵ The lineage predates GPT-2: the original GPT already tokenized with BPE, over a vocabulary of 40,000 merges, on top of a Transformer decoder.⁴ For this course we keep the character tokenizer — it makes every tensor small enough to print — but the embedding machinery below is identical whichever tokenizer feeds it. Only vocab_size changes.

The embedding table is a lookup

An integer id carries no useful geometry: token 43 is not "more" than token 20 in any sense the model should trust. The embedding layer fixes this by mapping each id to its own learned vector. In PyTorch, torch.nn.Embedding is documented as "a simple lookup table that stores embeddings of a fixed dictionary and size"; its two arguments are num_embeddings, the size of the vocabulary, and embedding_dim, the length of each vector. The input is a tensor of indices and the output is the corresponding vectors.⁸

Internally the module holds one weight matrix of shape (num_embeddings, embedding_dim). We will call the vector length d_model throughout the course, following the transformer paper's notation; it is the width every layer of the model works in.¹ Looking up a token is selecting its row.

import torch
import torch.nn as nn

vocab_size, d_model = 65, 384
token_embedding = nn.Embedding(vocab_size, d_model)
token_embedding.weight.shape        # (65, 384): one row per token

ids = torch.tensor([[20, 43, 50, 50, 53]])   # (B=1, T=5): one "Hello"
x = token_embedding(ids)                       # (B=1, T=5, d_model=384)
x.shape                                          # torch.Size([1, 5, 384])

Three axis names recur for the rest of the course, so fix them now: B is the batch (how many independent sequences we process at once), T is time or sequence length (how many tokens in each), and C is channels, our d_model (the width of each token's vector). The embedding step is the only place the last axis is created; from here to the output every layer maps (B, T, C) to (B, T, C).

It is worth seeing why a "lookup" is a real linear operation and not a special case. Selecting row i of a matrix is identical to multiplying that matrix by a one-hot vector with its 1 in position i. The embedding is that multiply, with the wasteful part skipped.

import torch.nn.functional as F

onehot = F.one_hot(ids, vocab_size).float()   # (B, T, vocab_size)
x_via_matmul = onehot @ token_embedding.weight # (B, T, d_model)

torch.allclose(x_via_matmul, token_embedding(ids))   # True

What the vectors learn

The rows of the embedding matrix are not features you compute; they are parameters you train. PyTorch initializes the weight from a unit normal distribution, and from there the vectors move under the same gradient descent that trains the rest of the network — there is no separate embedding objective.⁸ The geometry that results is a byproduct of the language-modeling loss: tokens that the model benefits from treating similarly drift toward similar vectors, because that is what lowers the loss.

The width d_model is a design choice, and it has grown with scale. The original transformer's base model used d_model = 512.¹ nanoGPT's default configuration reproduces GPT-2 small at n_embd = 768, with a vocabulary of 50,304 (GPT-2's 50,257 padded to a multiple of 64) and a block size of 1,024.⁷ A wider d_model gives each token more room to encode distinctions, at a quadratic-ish cost downstream; for our character-level model, 384 is comfortable.

Two refinements from the transformer paper are worth naming now, because they act directly on this table and you will meet them again in the code. First, the paper multiplies the embedding weights by the square root of d_model: "In the embedding layers, we multiply those weights by √d_model."¹ It is a scaling convention that keeps the embeddings on a comparable magnitude to the positional signal added next. Second, the same paper shares one weight matrix between the input embedding and the final layer that produces next-token scores: "we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation."¹ This weight tying, studied by Press and Wolf, roughly halves the embedding parameters and tends to lower perplexity, on the intuition that the space you read tokens from and the space you predict them into should be the same space.⁹

The shape of the input

Put the pieces in order. A batch of strings is encoded to a (B, T) integer tensor, the embedding table turns it into (B, T, C), and that tensor is what the first transformer block will read. The walk below is the forward pass up to the point the next lessons take over.

1. Encode

   Each string becomes a list of integer ids via stoi; stacked and padded to a common length, a batch is a (B, T) LongTensor of ids.⁶

2. Embed the tokens

   The token embedding maps (B, T) ids to (B, T, C) vectors — the operation of Fig. 3, a row lookup per position.⁸

3. Add positions (a preview)

   Token embeddings alone are order-blind: permute the sequence and the set of vectors is unchanged. GPT-style models add a second, position-indexed embedding of the same width, so x = wte(idx) + wpe(pos); the two tables have shapes (vocab_size, C) and (block_size, C).⁷ Lesson four builds positional encoding properly; here it is enough to know where it enters.

B, T, C = 4, 8, 384                          # batch, sequence, d_model
token_embedding = nn.Embedding(vocab_size, C)
position_embedding = nn.Embedding(block_size, C)  # block_size = max T

idx = torch.randint(0, vocab_size, (B, T))   # (B, T) ids
pos = torch.arange(T)                          # (T,) positions 0..T-1

tok = token_embedding(idx)                     # (B, T, C)
pos_emb = position_embedding(pos)              # (T, C), broadcast over B
x = tok + pos_emb                               # (B, T, C): ready for a block

Go deeper

Read Section 3.4, "Embeddings and Softmax", of "Attention Is All You Need" — it is four sentences and it states the embedding, the √d_model scaling, and the weight tying you now recognize.¹ Then open nanoGPT's model.py and find wte and wpe: the entire token-and-position construction of Fig. 5 is a handful of lines you can now read.⁷ To feel the tokenization trade-off in your hands, encode a paragraph of English with the character tokenizer above and with GPT-2's BPE, and compare the two sequence lengths against the 50,257-token vocabulary GPT-2 pays for that compression.³ If you want the algorithm itself, the merge loop in GPT-2's encoder.py is short enough to reimplement from scratch.⁵ And for the origin of subword segmentation, Sennrich et al. is the primary source.²