Implementing Byte Pair Encoding Tokenization

A token embedding is a vector representation of an input token, created in PyTorch with an Embedding layer.

Programming

Interactive lab

Try it: Text Vectorization: Tokens to Vectors

The text-to-tensor pipeline: standardize, split into tokens, build a frequency-ranked vocabulary with padding and [UNK], map tokens to integer ids, pad or truncate, then one-hot or multi-hot encode and look the ids up in an embedding table.

How it works

Standardize: lowercase and delete punctuation.
Split on whitespace into tokens.
Build the vocabulary: 0 = padding, 1 = [UNK], then words by frequency (ties: first appearance), capped at max tokens.
Map each token to its id (1 if it did not make the vocabulary), then pad with 0 or truncate to the sequence length.
Encode: one-hot rows per position or one multi-hot row per text; embed: id k selects row k of the embedding table.

Default run (18 steps): Raw text (59 characters): "The movie was great. The plot was great, the acting wasn't!". A model cannot read strings, so it is turned into numbers step by step. … Embedding lookup: id k selects row k of a fixed 8 x 3 table, giving a (8, 3) dense matrix instead of (8, 8) sparse one-hots. Row 0 (padding) is looked up too; a trained Embedding layer would learn these rows.

Simplified: Bounded ASCII text (up to 200 characters, 40 tokens), whitespace splitting only, and a small FIXED embedding table (a real Embedding layer learns its rows). The multi-hot vector keeps the padding slot (always 0) so its positions line up with the ids.

Educational simulation

Loading the simulation…

Predicting the Shape of an Embedding Output

Once a tokenizer has turned text into integers, the next question is what a model should do with those integers. In this lesson the token IDs come from the BPE tokenizer built earlier in the course, and each ID is passed through an embedding layer. The layer does not change the grid of IDs in a simple way. It replaces every integer with a vector, and the shape of the result tells you whether the layer behaves the way you expect.

What do you think happens?

A batch of 8 text samples, each with 4 token IDs, is passed into an embedding layer whose output dimension is 256. What shape does the output have?

(8, 4)
(8, 4, 256)
(256, 8, 4)
(32, 256)

Reveal answer

Answer:

(8, 4, 256)

The batch and sequence dimensions are kept as they are, and a third dimension of length 256 is appended. Every one of the 32 token IDs in the 8 by 4 grid becomes a 256-dimensional vector.

Keep your guess in mind. By the end of this article you should be able to derive that shape from the layer's constructor and the batch's grid, and to point to the exact step where a mismatch would first show up.

Why Token IDs Need to Become Vectors

A tokenizer turns text into integers. Each integer is a token ID that names one entry in a fixed vocabulary. An integer on its own carries no information about how one token relates to another, so a model cannot do much arithmetic with it directly. The usual remedy is to replace each ID with a vector: a list of numbers that the model can adjust during training. This replacement is called a token embedding.

The contrast between the two forms is the key idea of this lesson. A token ID is a single whole number that acts as a label. A token embedding is a dense list of floating-point numbers that the model can compare, combine, and tune.

Property	Token ID	Token embedding
Form	A single integer	A vector of floating-point numbers
What it names	One entry in the fixed vocabulary	The learned representation of that entry
Range or size	Integer between 0 and 50,256 in this lesson	A list of 256 values in this lesson
Can the model adjust it?	No, it is a fixed label	Yes, the values are learned during training
Relationship to other tokens	Carries none by itself	Can be compared and combined with other vectors

Token IDs and token embeddings carry different kinds of information.

Tracing One Token ID Through the Lookup

Think of the embedding layer as a table with one row for every possible token ID. The ID acts as a row index: when token ID 3 enters the layer, the layer returns row 3 of its weight matrix, and that row is the 256-value vector assigned to that token. Other IDs select their own rows, so each ID maps to its own vector.

When token ID 3 enters the layer, which row of the weight matrix is returned? Only the row matching the ID is selected.

The diagram shows only a few of the rows. The layer actually holds one row for each entry in the 50,257-entry vocabulary, and the highlighted row is the one the ID 3 selects. The other rows are untouched by that lookup.

Dimensions of the Embedding Weight Matrix

The layer is created with the constructor call torch.nn.Embedding(vocab_size, output_dim). In this lesson the first argument is 50,257 and the second is 256. These two numbers fully determine the shape of the layer.

Given vocab_size and output_dim, how many rows and columns does the embedding matrix have, and which number controls how many distinct tokens can be represented?

The vocabulary size sets how many rows the layer holds, one for each possible token ID. The output dimension sets how long each row is. A row is the vector that a given token ID maps to. So the number of distinct tokens the layer can represent is controlled by vocab_size, while output_dim controls how rich each token's vector is.

Following a Batch Through the Layer

The batch passed to the layer is a two-dimensional grid of token IDs, with one row per text sample and one column per token position. The lesson uses a data loader that fetches 8 samples with 4 tokens each, so the input grid has shape (8, 4). Every entry in that grid is one integer between 0 and 50,256.

If the input has shape (batch_size, seq_len) and the embedding dimension is emb_dim, what is the output shape at each stage?

Calling the layer on the grid replaces every integer with its 256-value vector. The grid positions do not move. The batch and sequence dimensions are preserved, and a third dimension of length 256 is appended. The result is a three-dimensional tensor of shape (8, 4, 256), so each of the 32 token IDs is now a 256-dimensional vector.

python

import torch

vocab_size = 50257
output_dim = 256
embedding_layer = torch.nn.Embedding(vocab_size, output_dim)

# A batch of 8 samples with 4 token IDs each
input_ids = torch.randint(0, vocab_size, (8, 4))
print(input_ids.shape)

embedded = embedding_layer(input_ids)
print(embedded.shape)

import torch

vocab_size = 50257
output_dim = 256
embedding_layer = torch.nn.Embedding(vocab_size, output_dim)

# A batch of 8 samples with 4 token IDs each
input_ids = torch.randint(0, vocab_size, (8, 4))
print(input_ids.shape)

embedded = embedding_layer(input_ids)
print(embedded.shape)

Output (expected)

torch.Size([8, 4])
torch.Size([8, 4, 256])

Notes on the code

The constructor arguments are vocabulary size first and output dimension second.
randint draws 8 by 4 random IDs from the range 0 up to, but not including, 50257.
The first printed shape is the input grid; the second is the embedding output.

Worked example

Predicting and checking the output shape

A batch of 2 samples with 6 token IDs each is passed into torch.nn.Embedding(50257, 256). What is the output shape, and how would you verify it?

1
Read the input shape
The grid has 2 rows (samples) and 6 columns (token positions), so the input shape is (2, 6).
2
Keep the batch and sequence dimensions
The embedding layer does not move grid positions, so the first two axes stay 2 and 6.
3
Append the embedding dimension
The output dimension is 256, so a third axis of length 256 is added.
4
Verify the result
Print the shape of the output and confirm the first two values match the input shape and the last value equals 256.

Result

The output shape is (2, 6, 256).

Where a Shape Mismatch First Appears

When you check embedding output, the source recommends a fixed order: verify the input shape first, then confirm that the final axis equals the output dimension. A mismatch is therefore detected at one of two places. If the input grid itself has an unexpected shape, the problem sits upstream of the layer. If the input is correct but the last axis of the output is not the output dimension, the layer was built with the wrong argument.

Common mistake
Using the vocabulary size as the output dimension
What it looks like
torch.nn.Embedding(256, 50257) is written when the intent was 50,257 rows and 256 columns.
Why it fails
The first argument sets how many rows exist, so the layer would accept only IDs up to 255 and return vectors of length 50,257.
Do this instead
Pass the vocabulary size first and the output dimension second: torch.nn.Embedding(50257, 256).
Common mistake
Expecting the output to lose the batch or sequence axes
What it looks like
Assuming an (8, 4) input becomes a flat (32, 256) matrix.
Why it fails
The layer keeps the batch and sequence dimensions and only appends the embedding dimension.
Do this instead
Expect (8, 4, 256) and check that the first two axes match the input grid.
Common mistake
Building two embedding layers with different sizes
Why it fails
The source notes that a second positional embedding layer must share the 256-dimensional size of the token embedding layer so that their outputs can be combined.
Do this instead
Use the same output dimension for every embedding layer whose outputs will be added together.
Common mistake
Checking only the numbers and ignoring the shape

Applying the Shape Check Yourself

EASY

A batch of 4 samples, each with 10 token IDs, is passed to torch.nn.Embedding(50257, 256). Before running anything, write down the input shape, the output shape, and the step where a mismatch would first show up if the output's last axis were not 256.

Hints

Start with the input grid: 4 rows and 10 columns.
The embedding layer keeps the first two axes and appends the output dimension.
Check the input shape first, then the final axis of the output.

Your attempt (not saved)

MEDIUM

Change the constructor to torch.nn.Embedding(50257, 128) and predict the new output shape for the same 4 by 10 batch. Then run the code and compare the printed shape with your prediction. Which axis changed, and why?