Tokens and embeddings

What you will be able to do

Give a network a useful way to read words. You turn text into token IDs and give each token a row of learned numbers. Then you check what those numbers learned, and what they did not.

NotebookTokens, embeddings and nearest words, in a notebookColab · Kaggle · GitHub

The notebook turns tickets into token IDs and looks up embedding rows. It trains the average-embedding network on the tickets (about one minute on a free Colab CPU) and finds the nearest words. Its checks tell you if your answer is right, without showing it. About 30 minutes.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 50 minutes

What you need to know

• The vocabulary and the bag of words, from The forward pass: every token that appears at least twice in the training tickets gets a position, and a ticket becomes a row of counts.
• Training with mini-batches, epochs and a validation set, from Batches, epochs and validation.
• Cosine similarity: how much two vectors point the same way, from 1 (same direction) to −1 (opposite), from Similarity and distance in the mathematics course.

What you need on your computer

• A browser. The notebook opens in Colab or Kaggle with PyTorch already installed.
• For the local route instead: the folder ticket-router with its active .venv, ticketnet.py and data/, from Set up PyTorch: Route 2.

All numbers come from the course’s recorded runs.

A list of names and a map

Look at an alphabetical list of towns. “Bath” comes right after “Barnsley”, but the two towns are far apart. “Bath” and “Bristol” are neighbours, but the list puts many names between them. The position in a list says nothing about where a town is.

A map is different. Each town gets two numbers, its coordinates. Towns that are close on the ground get close numbers, so you can measure how near they are. This lesson gives each word a “map position” of its own, and lets the network choose the positions.

From text to token IDs

A network computes with numbers, so the first step is always to cut the text into pieces and number them. The pieces are called tokens. This course uses the simplest tokens: lower-case words, numbers and single punctuation marks. The function tokenize in the course’s ticketnet.py does it with one regular expression.

Then each token gets its token ID: its row number in the vocabulary. You built the vocabulary in The forward pass: every token that appears at least twice in the training tickets. Larkfield’s vocabulary has 2,864 entries. ID 0 is <pad>, which fills up short tickets in a batch. ID 1 is <unk>, for any token that the vocabulary does not know.

tokens_and_ids.py

from ticketnet import Vocab, load_split, tokenize

train_rows = load_split("data/train.csv")
valid_rows = load_split("data/valid.csv")
vocab = Vocab([r["text"] for r in train_rows])

text = "Ugh, need to change trampoline delivery slot. Order LK-640933"
print(tokenize(text))
print(vocab.encode(text))
print(len(vocab), "entries;", vocab.itos[:2], "; ID of 'delivery':", vocab.stoi["delivery"])

The output:

['ugh', ',', 'need', 'to', 'change', 'trampoline', 'delivery', 'slot', '.', 'order', 'lk', '-', '640933']
[2653, 8, 1664, 2591, 618, 2620, 832, 2338, 10, 1749, 1486, 9, 1]
2864 entries; ['<pad>', '<unk>'] ; ID of 'delivery': 832

This is the real validation ticket T-64715, with 13 tokens. The order number 640933 appears in no training ticket, so it becomes ID 1, <unk>. About 1.2% of all validation tokens are <unk>, mostly order numbers like this one.

Why simple word tokens? Large language models cut text into smaller, frequent pieces, so that they never meet an unknown word. The lesson From text to tokens in the language models course shows how. Those tokenizers learn from billions of words. Larkfield has only about 4,700 training tickets. Word tokens are also easy to read: in Attention step by step, every token inside the network is a word you know. The cost is the <unk> token: a typo such as “delivry” loses its meaning completely.

The problem: an ID says nothing about meaning

Token IDs are like the alphabetical list of towns. “refund” and “refunded” get neighbouring IDs only because they are spelled alike. “parcel” and “courier” are far apart in the list, but they appear in the same kind of ticket. So the network cannot use the ID as an amount. ID 2,000 is not “more” than ID 1,000 in any useful sense.

Tomás tried it anyway: he gave a network the token IDs of each ticket as plain numbers. It routed only 29.0% of the validation tickets right, about as good as always answering “delivery” (27.5%). The failure case below has the details.

The bag of words of Module 2 avoids this. In effect, each token becomes a one-hot vector: 2,864 numbers, all 0 except a single 1 at the token’s ID. A ticket’s bag of words is the sum of its tokens’ one-hot vectors. That works, but one-hot vectors have two problems:

• They are long. 2,864 numbers for each token, and all but one of them are 0.
• Every two words are equally different. The one-hot vectors of “refund” and “refunded” share no position, so their dot product is 0 and their cosine similarity is 0. The same is true for “refund” and “trampoline”. The network must learn about each word separately, from the tickets that contain it.

The idea: a learned row of numbers for each token

Give each token a short row of numbers instead, for example 64 numbers. Put the rows in a table with one row per token ID: 2,864 rows × 64 numbers. To read a token, take its row. Then let training change the numbers, like any other weights.

The row of a token is called its embedding, and the table is the embedding table. The name comes from mathematics: you “embed” the words in a space of numbers, as a map places towns on a flat sheet with two coordinates. Here the space has 64 directions instead of two.

Why would the rows learn anything useful? Training changes each number so that the loss goes down. Words that push a ticket toward the same team, such as “charged” and “twice”, get rows that point in a similar direction. Similar rows lead to the same correct prediction. Nobody tells the network which words are similar. It finds the rows that help it route the tickets.

A tiny worked example. Take a vocabulary of four tokens and only 2 numbers per row:

ID	token	embedding row
0	<pad>	(0, 0)
1	parcel	(0.9, 0.1)
2	courier	(0.8, 0.3)
3	refund	(−0.2, 0.9)

The text “courier parcel” has the IDs 2 and 1. The lookup returns the rows (0.8, 0.3) and (0.9, 0.1), and their average is (0.85, 0.2). The cosine similarity of “parcel” and “courier” is (0.9 × 0.8 + 0.1 × 0.3) ÷ (0.906 × 0.854) = 0.97: they point almost the same way. For “parcel” and “refund”, it is (−0.18 + 0.09) ÷ (0.906 × 0.922) = −0.11. In a real network, training sets these numbers.

The lookup gives the same result as multiplying a one-hot vector by the table: the single 1 picks out one row. PyTorch skips the multiplication and reads the row directly, which is much faster. The diagram shows the whole path for the real ticket.

From text to embedding vectors: a table lookup. The IDs and numbers are real, from the course's trained average-embedding network.

Text description of the diagram

Step 1: the text of ticket T-64715, “Ugh, need to change trampoline delivery slot. Order LK-640933”. Step 2: thirteen tokens: ugh, comma, need, to, change, trampoline, delivery, slot, full stop, order, lk, hyphen, 640933. Step 3: their token IDs: 2653, 8, 1664, 2591, 618, 2620, 832, 2338, 10, 1749, 1486, 9 and 1. The order number is unknown, so it gets ID 1. Step 4: the embedding table has 2,864 rows of 64 learned numbers. Row 0, <pad>, is all zeros. Row 1, <unk>, starts −0.46, 1.36, 0.68, 0.92. Row 832, delivery, starts −0.32, −1.94, −0.63, 0.77. Row 2653, ugh, starts 1.67, −1.01, 3.08, 0.06. Step 5: the network reads the 13 rows, shape (1, 13, 64), averages them to one row, shape (1, 64), and an output layer turns it into 5 team scores. The answer is delivery, 99.3%.

In PyTorch, the embedding table is nn.Embedding(number_of_tokens, numbers_per_row). It starts with random numbers. padding_idx=0 keeps the <pad> row at zero and never trains it.

lookup.py

import torch
from torch import nn

torch.manual_seed(0)
embed = nn.Embedding(len(vocab), 64, padding_idx=0)   # 2,864 rows of 64 random numbers, to be learned
ids = torch.tensor([vocab.encode(text)])              # one ticket: shape (1, 13)
vectors = embed(ids)                                  # look up one row per token
print(ids.shape, "->", vectors.shape)
print(torch.equal(vectors[0, 6], embed.weight[832]))  # token 7 is 'delivery', ID 832
print(embed.weight.shape, embed.weight.numel(), "numbers to learn")

The output:

torch.Size([1, 13]) -> torch.Size([1, 13, 64])
True
torch.Size([2864, 64]) 183296 numbers to learn

Read the shapes. One ticket with 13 token IDs goes in. Out comes one row of 64 numbers for each token: (1, 13, 64). The new last dimension is the embedding. The vector of the seventh token is exactly row 832 of the table.

The average-embedding network on the tickets

Now route tickets with embeddings. The simplest network that uses them is the course’s AverageEmbeddingNet. It looks up the row of every token, takes the average of the rows, and sends the average through one output layer to 5 team scores. The average turns a ticket of any length into one row of 64 numbers. It skips the <pad> tokens, so padding does not change a ticket’s average.

train_average.py

from ticketnet import AverageEmbeddingNet, count_parameters, labels, set_seed, token_ids, train

ids_train, ids_valid = token_ids(train_rows, vocab), token_ids(valid_rows, vocab)
print(ids_train.shape, ids_valid.shape)
set_seed(0)
model = AverageEmbeddingNet(len(vocab))
print(count_parameters(model), "parameters")
history = train(model, ids_train, labels(train_rows), ids_valid, labels(valid_rows), epochs=30, lr=0.01, log=None)
best = min(history, key=lambda h: h["valid_loss"])
print("best:", best)
print("last:", history[-1])

The output:

torch.Size([4696, 120]) torch.Size([600, 120])
183621 parameters
best: {'epoch': 7, 'train_loss': 0.034, 'train_acc': 0.9934, 'valid_loss': 0.0704, 'valid_acc': 0.98}
last: {'epoch': 30, 'train_loss': 0.0009, 'train_acc': 1.0, 'valid_loss': 0.0959, 'valid_acc': 0.9683}

token_ids pads every ticket with zeros to the longest one, at most 120 tokens. The network has 183,621 parameters: 183,296 in the embedding table and 325 in the output layer (64 × 5 weights + 5 biases). At its best epoch, 7, it routes 98.0% of the validation tickets right. Then it overfits: after 30 epochs, 96.8%.

Compare this with the earlier routers. The bag-of-words network reached 98.3% at its best epoch, and the machine learning course’s linear model 98.8%. So the embeddings do not route better. That is no surprise: an average, like a count, throws the word order away. The next lesson shows what that costs.

What the embeddings learned

The embeddings are now trained, so you can ask: which words have rows that point the same way? Scale every row to length 1. Then the dot product of two rows is their cosine similarity. For one word, compute it with all 2,864 rows and keep the largest.

nearest.py

W = model.embed.weight.detach()
W = W / W.norm(dim=1, keepdim=True).clamp(min=1e-8)   # every row to length 1 (the <pad> row stays 0)

def nearest(word, k=5):
    sims = W @ W[vocab.stoi[word]]                    # cosine similarity with every word
    top = sims.topk(k + 1).indices.tolist()
    return [(vocab.itos[i], round(sims[i].item(), 2)) for i in top if vocab.itos[i] != word][:k]

for word in ("return", "refund", "cracked"):
    print(word, nearest(word))
print("refund / refunded:", round((W[vocab.stoi["refund"]] @ W[vocab.stoi["refunded"]]).item(), 2))

The output:

return [('collection', 0.76), ('sale', 0.75), ('back', 0.73), ('balcony', 0.73), ('returned', 0.7)]
refund [('twice', 0.57), ('bank', 0.56), ('paperwork', 0.56), ('payments', 0.55), ('statement', 0.54)]
cracked [('where', 0.61), ('along', 0.57), ('here', 0.56), ('lamp', 0.56), ('doorstep', 0.55)]
refund / refunded: -0.28

Read these lists honestly. They are partly meaningful and partly noise.

• “return” is close to “collection”, “back” and “returned”: the words of returns tickets. “sale” and “balcony” are there too, only because in these tickets they appear mostly in returns tickets (98% and 64% of the time).
• “refund” is close to payment words: “bank”, “payments”, “statement”. That fits Grace’s policy, where a confirmed refund that has not arrived is a payment problem. But it is not close to “refunded” (−0.28), which appears only 28 times in all the training tickets.
• “cracked” is close to “where”, “along” and “here”. These words share no meaning with it. In the training tickets, “cracked” appears in warranty tickets (51%) and in delivery tickets (46%, broken on arrival) almost equally. Its row cannot point to one team, so its neighbours are other words that point to no team.

The map shows the same thing as a picture. It starts from 43 chosen words: payment, delivery and warranty words, and some common words. Their 64 numbers are pressed down to 2 with PCA (principal component analysis), which keeps the two directions in which these rows differ most. The map shows 18 of the 43 words; the others sit so close together that their labels would overlap. Words that push toward the same team form groups. Common words such as “the” sit near the middle, because they push toward no team.

• mostly in delivery tickets
• mostly in payment tickets
• mostly in warranty tickets
• common words

The learned embeddings of 18 words, pressed from 64 numbers to 2 with PCA (computed on 43 chosen words). The two directions keep 51% of the differences between those rows, so the map is a flat shadow, not the whole picture. Each word has the colour of the team whose training tickets use it most.

Text description of the chart

A scatter plot of 18 words. The payment words charged and statement sit on the left, below the middle. The warranty words guarantee, stopped, repair and snapped sit at the bottom right. The delivery words tracking, arrived, box, porch, damaged, missing and late sit at the top. The common words the and my sit near the middle. Three words are in surprising places. Driver, a delivery word, sits at the top left, on the payment side of the map, far from the other delivery words. Cracked, a warranty word, sits on the right between the delivery and the warranty groups. Doorstep, a delivery word, sits at the right edge between the same two groups. The two plotted directions keep about half of the differences between the 43 chosen rows.

Look at the surprises too. “driver” appears mostly in delivery tickets, but its row sits on the payment side, and it pushes a ticket toward payment. “cracked” sits between the warranty and the delivery groups, and it pushes toward delivery. These rows learned what helped to route 4,696 tickets, nothing more.

Go deeper

The optional deep dive How embeddings learn meaning: word2vec and GloVe shows how word vectors can learn from the text alone, without team labels: the mathematics of skip-gram, GloVe and PMI.

When to use embeddings, and what they cost

Use an embedding table whenever a network reads tokens. Every network in the rest of this course starts with one, and so does every language model. Compared with a one-hot vector, an embedding is short (64 numbers instead of 2,864), and it can learn that two words act alike.

Know their limits:

• Similar is not the same meaning. Rows are close when words help the same prediction. “charged” and “twice” are close (0.75), but they do not mean the same thing.
• They reflect this data only. These rows learned from 4,696 tickets and one task, routing. They know nothing about English outside the tickets, and rare words get noisy rows.
• They cost parameters and data. The table has 183,296 numbers, almost all of this network’s weights. Each row learns only from the tickets that contain its word.
• An average forgets the order. Embeddings describe words, not sentences. How the words combine is the topic of the next two lessons.

Embeddings learned from a huge amount of text, in a pretrained model, know much more about English. Module 7 uses such a model.

Go deeper

The optional deep dive From counts to meaning compares one-hot counts, TF-IDF and dense embeddings with formulas, and measures why “parcel” and “package” only meet in an embedding.

Guided practice: nearest words in the notebook

1. Open the notebook and run the setup cells. They download the tickets and ticketnet.py.
2. Run the section “Tokens and IDs”. Check that ticket T-64715 has 13 tokens and that 640933 becomes ID 1.
3. Run the section “Train the average-embedding network”. It takes about one minute.
4. Run nearest("password") and nearest("charged").
5. For each list, write which neighbours make sense for Grace’s teams and which look like noise.

Expected output for step 4:

password [('associated', 0.57), ('cannot', 0.56), ('recognised', 0.56), ('service', 0.53), ('unable', 0.5)]
charged [('twice', 0.75), ('total', 0.74), ('charges', 0.71), ('identical', 0.68), ('bank', 0.67)]

Show an example answer

password: “cannot”, “recognised”, “unable” and “associated” come from sign-in problems: “my password is not recognised”, “I cannot log in”, “the email associated with my account”. They describe the account team’s tickets, not the meaning of “password”. “service” looks like noise.

charged: “twice”, “total”, “charges”, “identical” and “bank” all come from payment tickets about double charges and wrong totals. This list is the most meaningful one, because “charged” appears almost only in payment tickets (87% of the time).

The more a word belongs to one team, the more sensible its neighbours. Words that appear in several teams, such as “cracked”, get noisy neighbours.

Your turn: two words that look alike

Before you run anything, predict: is the cosine similarity of “refund” and “refunded” high, about 0, or negative? And of “cracked” and “broken”? Then use the notebook to check both.

Show the answer

“refund” and “refunded”: −0.28. “cracked” and “broken”: 0.49. A person sees two forms of one word in the first pair. The network sees two different patterns of use. “refund” appears 1,358 times, in tickets of all five teams. “refunded” appears only 28 times, so its row learned from very few tickets. Nothing in training pushed the two rows together. Embeddings measure “used alike for this task”, not “means the same”.

Failure case: token IDs used as numbers

Common mistake

Token IDs fed to the network as plain numbers

Symptom: Tomás skips the embedding table. He pads each ticket’s token IDs to 120 and divides them by 2,864, so they are between 0 and 1. Then he trains a network with one hidden layer (120 → 64 → 5). After 20 epochs, it routes only 29.0% of the validation tickets right. Always answering “delivery” gives 27.5%.

Cause: an ID is a name, not an amount. The network treats ID 2,000 as “twice as much” as ID 1,000. It also treats the first position of every ticket as one feature. The same word gets a different input in each position, and two unrelated words with close IDs look almost the same.

Fix: never feed IDs as numbers. Turn each ID into a vector first: a one-hot vector (the bag of words) or an embedding row (nn.Embedding). Then print the shape after the lookup: one row per token, (batch, tokens, 64).

See it end to endTrain a tiny language modelCase study

The tiny language model uses characters as tokens, and shows what that costs in sequence length.

Knowledge check

Check your understanding

In your own words: what is an embedding, and why is it more useful to a network than the token's ID?

Show the answer

An embedding is a row of learned numbers for one token, for example 64 numbers. The network finds it in the embedding table with the token’s ID, like a row number. The ID itself is only a name: a larger ID is not “more” of anything, and close IDs do not mean similar words. Training changes the embedding rows. So words that help the same prediction can get rows that point the same way, and the network can compare them.

Check your understanding

The trained embedding of 'cracked' has the nearest words 'where', 'along' and 'here'. Tomás says the embeddings are broken. What do you tell him?

Show the answer

The embeddings are not broken; they learned what routing needed from 4,696 tickets. “cracked” appears in warranty tickets and in delivery tickets almost equally, so its row cannot point to one team. Its neighbours are other words that point to no team. Embeddings show which words were used alike for this task, in this data. They are not a dictionary of meanings.

Check your understanding

Predict: an embedding table has 2,864 rows of 64 numbers. What shape does it return for a batch of 8 tickets padded to 120 tokens? How many numbers does the table hold?

Show the answer

The input is (8, 120) token IDs. The lookup adds one dimension of 64 numbers for every token, so the output is (8, 120, 64). The table holds 2,864 × 64 = 183,296 numbers. Training changes all of them, except the <pad> row when padding_idx=0 is set.

Recap

• Text becomes tokens (here: lower-case words, numbers, punctuation), and each token gets a token ID from the vocabulary of 2,864 entries. Unknown tokens, such as order numbers, become <unk>.
• An ID is a name, not an amount. A one-hot vector is 2,864 numbers long and makes every two words equally different.
• An embedding is a learned row of numbers for each token, found with a table lookup: (batch, tokens) → (batch, tokens, 64). Words that help the same prediction get similar rows.
• At its best epoch, the average-embedding network routes 98.0% of the validation tickets. That is no better than the bag of words (98.3%) or the linear model (98.8%). An average forgets the word order.
• Nearest words are partly meaningful (“charged” → “twice”, “total”) and partly noise (“cracked” → “where”, “along”). Embeddings reflect one task and 4,696 tickets.

Next, see exactly where counting and averaging words fail: tickets whose team depends on the order of the words.

Next: Why word order and context matter