Next-Token Prediction: Embeddings, and What Perplexity Means

Chapter 7 left you with a sequence of integers and no reason for one to follow another. Here is the reason, and it is one line of Chapter 2.

A language model is a function that takes the tokens so far and returns a distribution over which token comes next: one number per vocabulary entry, non-negative, summing to one. Nothing else. To get from that to a probability for a whole document, apply the chain rule of probability:

That is an identity, true of any sequence of anything, with no assumptions attached. So a model that does the small job — next token given the previous ones — has already done the large job of assigning a probability to every possible document, exactly and for free. The popular framing of this as a cheap trick ("it only predicts the next word") has the logic backwards: predicting the next token is modelling the joint distribution. There was never a second thing to do.

The loss follows just as mechanically. At each position the model produces a distribution $q$ and the truth is a single known token, so Chapter 4's cross-entropy applies unchanged:

That is the average negative log-likelihood — Chapter 2's recipe with a categorical distribution in the slot where the Gaussian sat. And since the true distribution is one-hot, its entropy is zero, so by Chapter 4's identity the cross-entropy equals the KL divergence: driving this number down and pulling the model's beliefs toward the data's are the same act.

One consequence deserves its own sentence, because it is the economic fact underneath the whole field. The labels are the data, shifted by one position. Nobody annotates anything. A trillion tokens of text is a trillion pre-labelled examples, which is why the training corpus of a modern model is "the internet" and not "a dataset somebody built".

Before any network, the baseline: 32,033 names, one per line, and the job of producing more of them one letter at a time.¹

The vocabulary is 26 letters plus a boundary symbol . marking both the start and the end of a name, so the model has to learn where names begin and where they stop. That is 27 symbols, and the smallest possible model is a table of how often each symbol followed each other symbol.

N = torch.zeros((27, 27), dtype=torch.int32)
for w in words:
    cs = ["."] + list(w) + ["."]
    for a, b in zip(cs, cs[1:]):
        N[stoi[a], stoi[b]] += 1

P = N.float()
P = P / P.sum(1, keepdim=True)            # one distribution per row   

Two lines of arithmetic and the model is fitted — and it is not a heuristic: dividing counts by row totals is the maximum-likelihood estimate for a categorical distribution, which is Chapter 2's recipe with the calculus already done.

names: 32033        train/val/test: 25626 / 3203 / 3204
training bigrams: 182583

the six most likely letters after 'a':
    a -> '.'  0.1944   a -> 'n'  0.1600   a -> 'r'  0.0967
    a -> 'l'  0.0749   a -> 'h'  0.0690   a -> 'y'  0.0606

Sample from it — pick a letter from the row of the current letter, move to that row, repeat until the boundary symbol comes up — and you get the names at the top of this chapter. They fail in a specific and informative way: locally plausible, globally nonsense. Every adjacent pair of letters in momakurailezitynn is a pair that occurs in real names; there are just seventeen of them in a row. The model has one letter of memory, so it cannot know it has been going on too long.

The loss on held-out names is 2.4546 nats. That number means nothing on its own, which is why perplexity exists:

Written out, with no library doing the work:

@torch.no_grad()
def perplexity(logits, Y):
    logp = F.log_softmax(logits, dim=1)          # log q for every symbol
    chosen = logp[torch.arange(len(Y)), Y]       # log q of the one that came next   
    return torch.exp(-chosen.mean())             

Exponentiating undoes the logarithm and returns the number to the units of counting things. The clean way to see what it counts is to measure a model that knows nothing at all — one that assigns probability $1/27$ to every symbol regardless of context:

uniform over 27 symbols            loss 3.2958 nats   ppl  27.000
bigram counts, add-one smoothed    loss 2.4546 nats   ppl  11.642

Exactly 27.000, because $e^{\log 27} = 27$. Perplexity is the effective number of equally likely options the model is choosing between. A perplexity of 27 means "no idea, could be anything". The count model's 11.642 means that one letter of context leaves it as uncertain as someone picking blindly from about twelve options instead of twenty-seven — which is why perplexity gets quoted and the raw loss does not.

Two things go wrong with it, and the second one goes wrong in published papers.

Zero probabilities are fatal. Of the 729 cells in the table, 113 never occur in training — 15.5 % of it is empty. That is fine until the held-out set lands in one, and seven bigrams in validation do, among them d→q, z→j and q→o twice. Probability zero means log $-\infty$, which means infinite loss and infinite perplexity: one name in three thousand destroys the metric. The usual patch is to add 1 to every count before normalising, which costs almost nothing here (2.4546 instead of 2.4524). But the patch is a confession. A count model cannot generalise at all. It has no way to suspect that q→o is plausible because q→u is common and o behaves like u elsewhere, since it has no notion that two symbols can resemble each other. Every cell is learned alone, and fixing that is what the rest of this chapter is for.

Perplexity is a price per token, and the token is a free parameter. This is the mistake that shows up constantly when models are compared, and it is easy to see once you look. Take the same corpus of English prose from Chapter 7, the same interpolated bigram model, and change only how the text is cut up:

Perplexity varies by a factor of 25 across those rows. Nothing about the model changed; only the size of the thing being predicted. Predicting a whole word is harder than predicting a letter, so it costs more per prediction — and there are fewer predictions to make.

Now read the last column, which divides the total cost by the number of characters instead and converts it to bits. It reorders the table. By perplexity the ranking is characters, words, BPE-512, BPE-2048; by bits per character it is words, BPE-2048, BPE-512, characters. The character model goes from first place to last. The 2,048-merge model, which by perplexity looks 6.6 times worse than the 512-merge one, is in fact the better of the two at 2.4254 bits against 2.6407.

So a perplexity is only comparable between two models that share a tokenizer, and models with different tokenizers can only be compared in bits per character — the quantity Shannon measured in 1951 by having human subjects guess the next letter of English text, and bounded at roughly one bit per character.² Our best bigram sits at 2.23 bits, which is a fair summary of how far this chapter still has to go.

Now build the same model as a network. It will take orders of magnitude more arithmetic to arrive at the same place, and arriving at the same place is the point.

Replace the table with one weight matrix $W$ of shape $27 \times 27$. Turn the current letter into a one-hot vector, multiply, and call the result logits — the unnormalised scores from Chapter 4. Then softmax, then cross-entropy, then gradient descent.

W = torch.randn((27, 27), requires_grad=True)

for step in range(3000):
    logits = W[xs]                            
    loss = F.cross_entropy(logits, ys)
    W.grad = None
    loss.backward()
    W.data -= 50.0 * W.grad

The highlighted line contains a definition worth having. Multiplying a one-hot vector by a matrix selects one row of it, so the multiply is a lookup — and every implementation skips the arithmetic and does the lookup directly, which is what W[xs] is.

That is an embedding table. A matrix with one row per vocabulary entry, indexed by token id. No geometry, no semantics, no separate algorithm: a lookup table whose contents happen to be learned by gradient descent along with everything else. Every mystical claim about "embedding space" bottoms out here.

Train it and watch where it goes:

  step     1   train 3.7550   val 3.3882   max gap to the count table 0.757269
  step   100   train 2.4732   val 2.4726   max gap to the count table 0.388354
  step  1000   train 2.4557   val 2.4549   max gap to the count table 0.041862
  step  3000   train 2.4547   val 2.4544   max gap to the count table 0.004048

The last column is the largest absolute difference between any cell of softmax(W) and the matching cell of the count table, and it goes to zero. After 3,000 steps the biggest disagreement anywhere in the 729 cells is 0.004048 and the mean is 0.000224. The worst cell is q→i, seen twelve times in the whole training set; among the 22 rows with more than a thousand occurrences the worst disagreement is 0.000562.

                 count table   network
    a -> '.'        0.1945     0.1945
    a -> 'n'        0.1601     0.1601
    a -> 'r'        0.0967     0.0967

Gradient descent, starting from random numbers and told nothing but "make the log-probability of the next letter large", rediscovered the table of counts. And it had to: the counts are the maximum-likelihood estimate, cross-entropy is the negative log-likelihood, so both procedures optimise the same objective and that objective has one optimum. The network did not learn something like counting. It converged to counting, slowly.

Which raises the fair question of why anyone would bother. Because the count table has nowhere to go from here, and the network does.

Extend the model to look at more than one previous character. This is Bengio's 2003 architecture, the direct ancestor of every model in the rest of this course:⁴ take the last three characters, map each through an embedding table into a 10-dimensional row, concatenate the rows into 30 numbers, push them through Chapter 5's hidden layer, and finish with an output layer producing one logit per vocabulary entry.

C  = torch.randn((27, 10))          # the embedding table
W1 = torch.randn((3 * 10, 200))     # the hidden layer from Chapter 5
W2 = torch.randn((200, 27))         # one output per vocabulary entry

emb = C[X].view(-1, 30)             # three lookups, concatenated   
h = torch.tanh(emb @ W1 + b1)
logits = h @ W2 + b2                
loss = F.cross_entropy(logits, Y)

Note what is new and what is not. The hidden layer is Chapter 5's, unchanged; the loss is Chapter 4's, unchanged. The novelties are the embedding table at the front and an output layer as wide as Chapter 7's vocabulary — and that second one is the expensive part of every language model ever built, because a real vocabulary has 100,000 entries and this matrix multiply runs at every position.

The same code, trained identically, with only the size of the context window changed:

The second row is the interesting one. A network with a 200-unit hidden layer and eleven times as many parameters as the count table performs exactly as well as the count table and no better. Capacity was never the limitation. One character of context permits a certain loss and nothing you bolt on can go below it, because the information is not there.

Give it three characters and the perplexity drops from 11.68 to 8.29 — a 29 % cut, bought with 4,000 extra parameters. It beats counting here for precisely the reason diagnosed earlier: a count model over three-character contexts needs $27^3 = 19{,}683$ rows, most of them empty or holding a single observation, and it learns every one alone. The network shares. If a, e and i end up with similar embedding rows, what it learns after bra transfers to bre without it ever having seen bre. That transfer is the whole value of the embedding table, and it is the gap between rows two and three.

The samples improve accordingly:

deliah   nellara   joce     kael      quintis
salayson  reety    khyrmin  mahnen    madiaryxia

Still not a list of real names. But deliah, nellara and kael would not look out of place on one, and the run-on monsters are gone: the longest of twenty samples from the count model is nineteen letters, the longest of twenty from this one is thirteen.

The table is $27 \times 10$: one row of ten numbers per character, all initialised randomly and moved only by the gradient of the next-character loss. Nobody put anything in there. So what ended up in it?

The tool for asking is cosine similarity, which is the dot product of Chapter 1 with the lengths divided out:

It measures the angle between two vectors and ignores their lengths, which is what you want when a row's length reflects how often its token appeared rather than what it means. Normalise every vector to length 1 first — as real systems do, once, at indexing time — and cosine similarity is simply the dot product.

Here are the nearest neighbours of a few characters in the trained table:

  'c' -> 'k':+0.598      'j' -> 'z':+0.650      'i' -> 'y':+0.541
  'u' -> 'e':+0.482      'a' -> 'h':+0.367      '.' -> 'q':+0.077

Some of that is what the folklore promises. c and k are interchangeable in names, and so are i and y; j and z are both rare, mostly-initial consonants that behave alike. The boundary symbol . is near nothing at all — 0.077 to its closest letter — because it is the only symbol that marks a position rather than a sound.

And some of it is not. The nearest neighbour of a is h, not another vowel. Averaged over all pairs:

mean cosine, vowel to vowel         : +0.1889
mean cosine, consonant to consonant : +0.0765
mean cosine, vowel to consonant     : -0.0042

The vowels are more like each other than like consonants, and the effect is real but small. Tested against 2,000 randomly chosen groups of five letters, 58 of those groups separate at least as cleanly — a gap significant at about $p = 0.03$. Real, then, and nothing like the crisp geometric island that popular accounts of embeddings imply.

That is the honest description of an embedding table and it is worth holding on to for the rest of the course. It is not a map of meaning. It is a change of coordinates, learned rather than designed, whose only job is to make the next layer's job easy — the same sentence Chapter 5 used for the hidden layer that folded the plane to solve XOR. Any structure you find in it is there because it lowered the loss, and structure that does not lower the loss is simply not there.

If the useful part is the table, you can go after it directly. That is word2vec: keep the embedding lookup, throw away the language model.⁵

The skip-gram with negative sampling objective is one line. For a real (centre, context) pair drawn from the corpus, push their dot product up; for $k$ fake pairs drawn from a noise distribution, push it down:⁶

That is a binary classification — "did these two words really occur together?" — and it is cheap precisely because it never touches the full vocabulary, which is what made training on billions of words practical in 2013. GloVe arrives at similar vectors from the other direction, by factorising the matrix of global co-occurrence counts instead of streaming through examples.⁷ Both are fitted to exactly the statistic the count table was built from. They are counting, compressed.

Trained on text8 — 17,005,207 words of English Wikipedia, 71,290 of them occurring at least five times, 100 dimensions, three passes — the vectors come out with the property that made them famous:

king     -> charles 0.700, son 0.693, queen 0.686, henry 0.669, throne 0.667
physics  -> chemistry 0.672, electromagnetism 0.661, quantum 0.654, theoretical 0.624
guitar   -> bass 0.733, vocals 0.732, acoustic 0.728, guitars 0.703, drums 0.685
three    -> seven 0.892, two 0.877, one 0.875, five 0.871, four 0.870

Nobody supplied a category for instruments or for numerals. Now the famous part: take king, subtract man, add woman, and find the nearest vector to the result.

king - man + woman
   nothing excluded : king 0.693, elizabeth 0.657, wife 0.629, woman 0.607
   a, b, c excluded : elizabeth 0.657, wife 0.629, mary 0.607   (queen is 4th, 0.604)

The nearest vector to king - man + woman is king. That is not a quirk of one example. Mikolov's evaluation set poses questions of the form a : b :: c : ? — 8,869 semantic ones (paris : france :: rome : italy) and 10,675 syntactic ones (walking : walked :: swimming : swam) — and across the 4,103 semantic questions this vocabulary can answer, the winner is one of the three input words 99.8 % of the time. The published demonstrations do not mention it, because the standard scoring rule deletes a, b and c before looking. It is a legitimate rule, and it is doing more work than the arithmetic:

The third row is the one to sit with. Throw away a and b, do no arithmetic at all, return whatever is nearest to c — and you keep 77 % of the semantic score. Most of what looks like analogical reasoning is proximity plus a rule that forbids the obvious answers, which is what Linzen measured on properly trained vectors and what the baselines above replicate.⁸ These particular vectors are small — 17 million words against the billions behind the published models — so read the percentages as a shape, not a state of the art. The shape is what survives at every scale: the arithmetic is real, and far weaker than the one demonstration everybody quotes.

Everything so far has a hard limit built into the data structure. A table has one row per token. The word bank gets one vector, the same one in a sentence about a river and a sentence about a mortgage — necessarily, because a lookup by id cannot depend on anything else.

The fix is to stop reading the vector out of the table and start computing it from the sentence. That is a contextual embedding, introduced by ELMo in 2018 and made standard by BERT the same year.⁹¹⁰ Measured on the real model, the numbers are sharper than the explanation:

sentence A: "He sat on the bank of the river and watched the water go by."
sentence B: "She deposited the cheque at the bank on the corner of the street."

static vector for 'bank' (a row of the input embedding table)
    cosine A vs B ........................ 1.000000

contextual vector for 'bank', layer by layer
    layer  |  A vs B  |  A vs another river sentence  |  B vs another money sentence
        0  |  0.9512  |            0.9512             |            0.9359
        4  |  0.5647  |            0.8987             |            0.7716
        9  |  0.4284  |            0.8699             |            0.7568
       12  |  0.5278  |            0.8702             |            0.7335

The first row is exact, not approximate: the static vector for bank is the same 768 numbers in both sentences, so the cosine is 1 by construction. Nine layers later the two occurrences sit at 0.43, while bank in two different river sentences stays at 0.87. Nobody labelled a sense anywhere in this process; the senses separated because separating them makes the training objective — guessing a hidden token from its neighbours — easier to satisfy.

Two details repay attention. Layer 0 is already 0.9512 rather than 1.0, because position embeddings have been added and the word sits in a different place in each sentence. And the similarity rises again at layers 11 and 12: the final layers of a pretrained model are specialised to its training objective, and are often not the best place to take a representation from.

To search a corpus by meaning you need one vector per sentence. Given those, the search is trivial — this is the whole of semantic retrieval, and Chapter 19 is about everything around it:

E = normalise(embed(sentences))       # (200, d), every row of length 1
q = normalise(embed([query]))         # (1, d)
scores = q @ E.T                      # one matrix multiply   
top5 = scores[0].argsort()[::-1][:5]

So the only real question is where embed comes from. The obvious move is to take a pretrained language model, run each sentence through it and average the token vectors. Here is that method against four alternatives, scored two ways: the rank correlation between cosine and human similarity judgements over the 1,379 pairs of the STS benchmark, and top-1 retrieval on an index built from the 200 most strongly paraphrased of those pairs — one side of each pair indexed, the other used as the query.

Read the middle three rows against the first two. A 109-million-parameter pretrained transformer, used the obvious way, is worse at judging sentence similarity than counting how many words two sentences share — and worse than averaging the 100-dimensional text8 vectors trained a moment ago. The [CLS] token, which tutorials still recommend because BERT was pretrained with a sentence-level objective attached to it, is worse than half of that.

This is not a defect in BERT. It is the objective. A language model is trained so that its hidden states predict a token; nothing there asks two paraphrases to end up near each other, and nothing rewards a geometry in which cosine means "same meaning". The last row is a model a fifth of the size (22,713,216 parameters) trained on a different loss entirely: contrastive learning, where the examples are pairs — a question and its answer, a sentence and its paraphrase — and the objective pulls true pairs together while pushing sampled negatives apart. That is Sentence-BERT's contribution and the origin of the whole embedding-model industry.¹² Dense Passage Retrieval applies the same recipe to search directly, with one encoder for queries and one for passages.¹³

So, the practical rule:

Chapter 19 builds a production retrieval system on this footing and arrives at a concrete cosine cut-off. The last measurement in this chapter is what makes such a number defensible rather than magic.

Real embeddings have hundreds or thousands of components, and distances behave strangely up there. Take 1,000 random points in the unit cube of $d$ dimensions and look at the ratio between the largest and the smallest distance between any two of them:

In ten thousand dimensions the farthest pair of points is only 6 % further apart than the closest pair. Everything is roughly equidistant from everything else, "nearest neighbour" stops carrying much information, and that is the curse of dimensionality — as well as one reason large vector databases do not do exact nearest-neighbour search. The other side of the same coin is what makes cosine thresholds workable: measured over a thousand pairs of random unit vectors, the mean cosine sits at $-0.0052$ in 100 dimensions and $+0.0003$ in 768, with standard deviations of 0.0968 and 0.0357 — and in 768 dimensions only 0.2 % of random pairs exceed 0.1 in absolute value. A measured similarity of 0.4 is therefore not "40 % alike"; it is far outside anything chance produces, which is why thresholds between 0.3 and 0.7 separate signal from noise instead of sitting in the middle of it.

The model in this chapter reads a fixed number of previous characters, looks each one up and glues the results together in order. That design has two problems, and they are the same problem.

Look again at the context table: going from three characters to eight nearly doubled the parameters and bought 0.06 nats. The cost grows linearly with the context — every extra position needs its own slab of the first weight matrix — and the benefit does not. Push it to a thousand tokens and the first layer alone outweighs the rest of the model, most of it spent on positions that do not matter for any given prediction.

Which is the second problem: the model has no way to decide which of the previous tokens matter. Position two gets its own weights and position seven gets its own, permanently, whatever is in them. When the model is spelling nell, the decisive character is the one immediately before. When a sentence contains a pronoun, the word that fixes its referent may be forty tokens back — and no fixed slot can be assigned to "forty back", because next time it will be six.

What we want is a model that computes, for each prediction, how much each earlier token should count — weights over the context produced by the content rather than fixed by the layout. Write that down carefully and it begins as something entirely mundane: an average over the previous tokens. Then let the weights of that average be learned, and let them depend on which token is doing the asking.

That is attention, and it is Chapter 9.

Also worth reading alongside: chapter 3 of Jurafsky and Martin's Speech and Language Processing, which treats n-gram models, smoothing and perplexity far more carefully than there is room for here, including why interpolation and back-off beat adding one; the Stanford CS229 notes §17.1–17.2 for language modelling from the probabilistic side; and Linzen's paper above, which is short and worth reading in full.