Data Loaders for Next-Token Prediction

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…

Why a Token ID Is Not Enough

A model that predicts the next token in a text has to read tokens as input first. Before any learning happens, text is split into tokens by a tokenizer, and each token is mapped to an integer. In this course the tokenizer is the BPE tokenizer built earlier, and its vocabulary has 50,257 entries. Each integer is a token ID that names exactly one entry in that fixed vocabulary.

A single integer carries no information about how one token relates to another. The ID 42 is not 'close to' 43 in any meaningful sense; it is just a label. Because of this, a model cannot do much arithmetic with the raw integer. The usual remedy is to replace each ID with a vector, a list of numbers that the model can adjust during training. That replacement is called a token embedding, and the layer that performs it is an Embedding layer in PyTorch.

Before reading the rest of this article, try to predict what the Embedding layer returns for a batch of token IDs. Think about what the output should look like: will it still be a grid of integers, or something with more structure? The answer is below, and it is worth committing to a guess first.

Tracing One Token ID Through the Layer

The Embedding layer can be understood as a table with one row per possible token ID. Looking up a token ID means reading the row whose index equals that ID, and the row returned is the vector for that token. The ID does not change the table; it only selects which row to return.

When the token ID 42 enters the Embedding layer, the layer returns the row at index 42 of its weight table, and that row is the 256-value vector for that one token.

Notice what the diagram shows and what it does not. The neighboring rows 41 and 43 are not involved in the lookup for ID 42. Each ID maps to one specific vector, and different IDs select different rows. The vectors are not computed from the integer's numeric value; the integer only serves as an index into the table.

python

import torch

embedding = torch.nn.Embedding(50257, 256)

token_id = torch.tensor([42])
vector = embedding(token_id)

print(vector.shape)

import torch

embedding = torch.nn.Embedding(50257, 256)

token_id = torch.tensor([42])
vector = embedding(token_id)

print(vector.shape)

Output (expected)

torch.Size([1, 256])

Notes on the code

torch.nn.Embedding(50257, 256) creates a table with 50,257 rows, each of length 256
token_id holds a single ID, placed in a one-element tensor
Calling the layer returns the row for that ID, so the result has one vector of 256 values

Reading the Constructor Arguments

The layer is created with the form 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 layer's shape. The first argument sets how many rows the layer holds, one for each possible token ID. The second sets how long each row is, which is the length of the vector a token ID maps to.

Constructor argument	Value in this lesson	What it controls
vocab_size (num_embeddings)	50,257	Number of rows, one per possible token ID
output_dim (embedding_dim)	256	Length of each row, the vector size per token

The two constructor arguments and the parts of the weight table they set.

Given vocab_size and output_dim, the weight matrix has one axis for the number of tokens and one axis for the number of vector components.

Read the matrix as a table with a row for every token and a column for every component of that token's vector. The first axis counts tokens, so it is set by vocab_size. The second axis counts vector components, so it is set by output_dim. Swapping the two arguments would produce a completely different layer, one that holds 256 rows and returns vectors of length 50,257, so the order of the arguments matters.

How a Batch Changes Shape Through the Layer

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

What do you think happens?

The input grid has shape (8, 4). What shape do you expect the Embedding layer to return?

(8, 4), the same as the input
(8, 4, 256), with a new axis of length 256 appended
(256, 8, 4), with the embedding axis moved to the front

Reveal answer

Answer:

(8, 4, 256), with a new axis of length 256 appended

The layer keeps the batch and sequence dimensions and appends the embedding dimension. Each of the 32 token IDs becomes a 256-value vector, so the grid becomes a three-dimensional tensor.

The grid of token IDs keeps its batch and sequence dimensions, and the layer appends the 256-length embedding axis to each position.

Worked example

Predicting the output shape for a batch of 8 by 4 token IDs

A batch of token IDs has shape (8, 4). The Embedding layer is torch.nn.Embedding(50257, 256). What is the shape of the output, and how many 256-value vectors does it contain?

1
Check the input shape
The input is $(8, 4)$: 8 samples, each with 4 token positions, giving 32 IDs in total.
2
Identify the embedding dimension
The second constructor argument is 256, so every ID is replaced by a vector of length 256.
3
Keep the batch and sequence axes
The layer does not move grid positions, so the first two axes stay at 8 and 4.
4
Append the embedding axis
A third axis of length 256 is added, giving $(8, 4, 256)$.
5
Count the vectors
Each of the 32 positions holds one 256-value vector, so there are 32 vectors.

Result

The output has shape $(8, 4, 256)$ and contains 32 vectors of length 256.

python

import torch

embedding = torch.nn.Embedding(50257, 256)

# Stand-in for a batch from the data loader: 8 samples, 4 tokens each
token_ids = torch.randint(0, 50257, (8, 4))

print(token_ids.shape)

embedded = embedding(token_ids)
print(embedded.shape)
print(embedded.shape[-1] == 256)

import torch

embedding = torch.nn.Embedding(50257, 256)

# Stand-in for a batch from the data loader: 8 samples, 4 tokens each
token_ids = torch.randint(0, 50257, (8, 4))

print(token_ids.shape)

embedded = embedding(token_ids)
print(embedded.shape)
print(embedded.shape[-1] == 256)

Output (expected)

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

Notes on the code

torch.randint(0, 50257, (8, 4)) creates random IDs in the valid range, standing in for real batch contents
The first print confirms the input grid is 8 by 4
The second print confirms the output has the appended 256-length axis
The last check tests the final axis, which must equal the output dimension

Where Shape Mismatches Show Up First

When you check embedding output, the source recommends a fixed order of checks. First verify the input shape. Then confirm that the final axis of the output equals the output dimension. Checking in this order tells you exactly where a problem starts: if the input shape is already wrong, the error originates upstream, in how the batch was built; if the input is correct but the last axis is not 256, the error is inside the embedding setup.

Common mistake
Expecting the output to keep the input's shape
What it looks like
Assuming a (8, 4) input returns (8, 4) because it is a grid of IDs
Why it fails
The layer replaces each integer with a vector, so a new embedding axis is appended.
Do this instead
Expect $(8, 4, 256)$: the first two axes are kept and a length-256 axis is added.
Common mistake
Swapping the two constructor arguments
What it looks like
Writing torch.nn.Embedding(256, 50257) and expecting the same layer
Why it fails
The first argument sets the number of rows, one per token ID, and the second sets the length of each row.
Do this instead
```
Use torch.nn.Embedding(vocab_size, output_dim), here torch.nn.Embedding(50257, 256).
```
Common mistake
Checking only the printed numbers and not the shape
Why it fails
The source emphasizes that the shape is where you can verify the layer behaves as expected; plausible-looking values can hide a wrong axis.
Do this instead
Print the shape before and after the layer, and compare the final axis to the output dimension.
Common mistake
Treating the embedding as a function of the integer's numeric value rather than a row lookup

Practice: Predict and Check a Shape

EASY

A batch of token IDs has shape (4, 6). Using torch.nn.Embedding(50257, 256), predict the shape of the output. Then state which axis you would check first if the output shape turned out to be (4, 6) instead of the value you predicted.

Hints

The embedding layer keeps the batch and sequence axes, so the first two numbers of the output match the input.
The output dimension is the second constructor argument, which is the length of the appended axis.
Check the input shape first, then the final axis of the output.

Your attempt (not saved)

MEDIUM

Write a short script that builds an Embedding layer with 50,257 rows and 256 columns, passes in a random batch of shape (2, 3), and prints the output shape. Then change the second constructor argument to 128 and predict, before running it, how the final axis of the output changes.