Words as vectors: tokens and embeddings

The problem: getting a computer to read

Read "Once upon a time, there was a little" and you already expect "girl" or "dog", never "upon". A computer gets only the characters, each stored as a number that says nothing about meaning. Guessing what comes next in a text is called language modeling. Phone keyboards do it to suggest your next word, and the large language models behind chat assistants are, at their core, very good guessers of the next piece of text; chapter 8 trains a tiny one.

Why guess the next word?

Because guessing well takes a lot of knowledge: grammar ("a girl", not "a girls"), meaning (after "she ate a" comes something to eat), and the story so far. And text is its own answer key. Nobody has to label anything, because every word is the right answer for the text before it. This chapter's network learns from 20,000 short stories from TinyStories, a collection of stories written to use only words a 3- or 4-year-old understands. Researchers at Microsoft had GPT-3.5 and GPT-4 write them in 2023, to study how small language models learn to write; this chapter uses stories written by GPT-4.

Why is it hard?

Words are symbols, not numbers. "Cat" and "dog" share no letters but are used in the same places; "cat" and "car" share most of their letters but are not. Worse, much of any text is new, and the longer the stretch, the more of it is new. Here is a story the network never saw, cut into its tokens (the next section explains them). A token is marked when the tokens just before it, in that order, appear nowhere in the 20,000 training stories.

The simplest idea is to count: after the last few tokens, which token came next most often in the training stories? When those tokens never appeared together, fall back to fewer of them. Counting works best when the exact phrase was seen; when it was not, it can only answer from a shorter piece of it. Looking further back does not help: long stretches rarely repeat, so counting with the last 8 tokens does no better than with the last 4.

What a neural network does differently

• A token becomes a list of numbers. Each of the 2,000 tokens gets its own list of 32 numbers, called its embedding. The network reads those numbers, never the spelling.
• Words used alike get similar lists. The lists are learned, and tokens that appear in similar places end up with similar numbers. What the network learns about "cat" then carries over to "dog", and a phrase it never saw can still look like phrases it did.
• It learns from its mistakes. The network is trained on one task, guessing the next token, by gradient descent on the 4.3 million tokens of the training stories.

Measuring…

What this is

This film follows one phrase from characters to the network's guess, then shows what training did to the lists of numbers.

The question

Counting can only answer for phrases it has seen. What does a network do with a phrase its training stories never contain? Pick one, or type your own, and compare.

Counting needs the exact tokens. When the last 8 have never appeared together, it falls back to fewer, until it finds a stretch it has seen, and then answers from that alone, forgetting the rest. The network always reads all 8 tokens, through their lists of numbers, and usually gives a sensible guess even for a phrase that appears nowhere in its training. The tools below show how.

Tool 1: cutting text into tokens

A network reads from a fixed list of tokens, its vocabulary. A token for every whole word would need a very long list, and new words would still turn up: the 20,000 training stories use 9,597 different words (counting forms such as "lily's" separately), and 1,000 other stories bring 99 more, such as "tulip", "pudding" and "excellent". Single letters need only a few dozen tokens, but every word becomes a long row of them, and a letter alone says little. Most language models use pieces in between, and a common way to find them is byte-pair encoding: start from single letters, find the pair of neighbors that appears most often in the training text, glue it into one new piece, and repeat. Drag the slider to replay the 1,922 merges learned from the training stories.

After all 1,922 merges the vocabulary has 2,000 tokens: the end of a story, a line break, 13 punctuation marks, the 10 digits, 52 single letters (a letter that starts a word, or one glued on with "##"), a glued-on apostrophe, and the 1,922 pieces. Common words become one token, as 93% of the words in other stories do, counting every time a word appears; rarer words are cut into a few pieces, such as "tow" and "##ers"; and any word at all can still be spelled out letter by letter. The tokenizers of GPT models work the same way on bytes instead of letters, keep capitals and spaces inside their tokens, and have between about 50,000 and 200,000 tokens.

Tool 2: a token becomes a list of numbers

Each token is now a number, its place in the vocabulary, and that number means nothing: token 1,203 is not "more" than token 87. So every token gets a list of 32 numbers of its own, its row in a table of 2,000 rows, the embedding table. Reading a token means looking up its row. Here are some rows of the finished network's table; type any word to add its row.

Looking up a row is a layer of weighted sums in disguise, like those of chapter 1 but without the bias or the tanh. Write the token as 2,000 numbers that are all 0 except for a single 1 at its place, a one-hot list, and multiply it by the table: only the row at the 1 survives. One-hot lists say nothing about similarity: every two of them are exactly as far apart, so "cat" differs from "dog" as much as from "the". The table is where similarity can live, and training is what puts it there.

Tool 3: the network that guesses the next token

The network reads the last 8 tokens. It looks up their 8 rows and lays them side by side into one list of 256 numbers. A hidden layer of 128 tanh neurons, as in chapter 1, reads that list, and an output layer gives each of the 2,000 tokens a score. Softmax turns the scores into probabilities. Type a phrase and follow it through.

Training is the backward pass of chapters 3 and 4. The loss is −log of the probability the network gave to the token that really came next, and its gradient flows back through the output layer and the hidden layer into the 8 rows that were read. Those rows, and only those, get a gradient from this example, and it nudges them toward what would have helped; the table is just more weights, learned like the others. One change from earlier chapters: the steps are taken with Adam, a version of gradient descent that gives every weight its own step size. The row of "the" gets a gradient in almost every batch, the row of a rare word once in a long while; Adam divides each weight's step by the typical size of its recent gradients, so both move at a useful pace. Most language models are trained with Adam or a close relative.

Secret 1: nobody told it what words mean

The table starts as small random numbers, so at first every word is nowhere in particular. The map below shows 42 words in six groups. A row has 32 numbers, too many to draw, so the map shows the flat view of the rows that keeps the most of their spread (principal component analysis); words with similar rows land close together on the map (two words can also land close without having similar rows, because the flat view keeps only part of the spread). The early moments are that network's own first steps, replayed in your browser right now: a fresh network with the same starting numbers, trained on the same batches. The later moments were saved by the run that trained the finished network.

Nobody told the network that "red" and "blue" are colors or that "park" and "garden" are places. It only ever tried to guess the next token, and a token's row changes only when that token is among the 8 being read, to help guess what comes after them. So words that are followed by similar tokens got pushed in similar directions. Colors gather because they come before the same words, most often "ball" and "car". The actions gather less, and the map shows why: "ran" comes before "to" or "away", like other ways of moving, while "laughed", "smiled" and "cried" come before "and" or a full stop, and sit near the feelings.

Secret 2: close means used alike, not the same

Pick a word to see the words whose rows are closest to its row. Closeness is measured by the angle between two rows (the cosine similarity: 1 when they point the same way, 0 when they are at right angles, negative when they point apart).

Opposites are often neighbors: "big" and "small", "up" and "down". They fit the same slots ("a … dog", "climbed … the hill"), so guessing the next token after either is much the same job. The table holds how words are used. That is related to what they mean, but it is not the same thing.

Secret 3: directions mean something too

The steps between rows can carry meaning as well. If the step from "he" to "she" is roughly the same as the step from "boy" to "girl", then the row of "she", minus the row of "he", plus the row of "boy", should land near the row of "girl". Try it.

Many of these puzzles come out right, but not all, and the result often depends on a detail: whether the three words of the puzzle may themselves be the answer. Left in, the nearest word is often one of them, usually the one added last; famous examples such as king − man + woman ≈ queen are reported with them left out. Those come from much larger networks trained on far more text. ("Queen" and "woman" are too rare in these stories to be single tokens here.)

Secret 4: one word, one row

Each token has a single row, whatever it means in its sentence. "Saw" is a tool in "he cut the wood with a saw" and a past tense in "she saw a bird", and both read the same 32 numbers. Pick a word with two meanings to see it used both ways in the training stories, and what the network expects next in each.

The network can still tell the two meanings apart a little, from the 7 tokens before it, but the word's own numbers cannot change with its sentence, and its row ends up mostly about the more common meaning. The network also has a short memory: it sees 8 tokens back, so a name mentioned 30 tokens ago is gone. Both are what attention fixes, in chapter 7: it rebuilds each token's numbers from the tokens before it, as far back as the model's window reaches.

The algorithm: tokens, and the network

Seven parts: three that cut text into tokens, and four that make up the network and its training. This is the JavaScript this page actually runs, or the same program in Python.

The playground

Start a story and watch it being cut into tokens and what the network expects next. Let it write on, look up any word's neighbors, or train your own network from scratch.