insideout Get Pro

09 · Logits and softmax over the vocabulary

Pro

Transformers, lesson 9

Logits and softmax over the vocabulary

The last layer scores every token, and softmax turns scores into probabilities.

You need "Norm, attention, MLP, residual" and "Looking up a token's vector". This lesson picks up where the last block hands over its vectors.

Part 1. From one vector to 50,257 scores

Feed GPT-2 small the text "The cat sat on the". After 12 blocks and a final layer norm, each of the 5 positions holds 768 numbers. To guess the next token, the model takes the vector at the last position and multiplies it by one more weight grid, 768 × 50,257. Out come 50,257 numbers, one per entry in its vocabulary.

Example 1.1

The last step of a forward pass

  1. "the" (last position): 768 numbers

    Only the last position's vector is used to guess what comes next. It has gathered context from all the earlier tokens through attention.

  2. 768 numbers×grid 768 × 50,257=50,257 scores

    One weighted sum per vocabulary entry. Each column of the grid acts as a pattern for one token, and the score says how well the vector matches it.

  3. mat 3.2floor 2.5couch 2.0roof 0.2... 50,253 more

    Higher means "more likely next". These raw scores can be any size and any sign, and they don't add up to anything in particular. (Scores here are illustrative.)

Logits

The raw scores, one per vocabulary token, before they're turned into probabilities. The grid that produces them is the output layer (also called the unembedding or LM head).

Your turn

Llama 2 7B has a vocabulary of 32,000 tokens and vectors of 4,096 numbers. How many logits does it produce for each next-token guess?

One score per vocabulary token, so 32,000. The 4,096 is the size of the vector going in. 4,096 × 32,000 is the size of the output grid, about 131 million weights. The model never outputs a single ID; something outside it has to pick.

Part 2. Softmax turns scores into probabilities

Keep going with Pro

You've done the first part. Pro opens the rest of this lesson, and every other one.

Still to come
  1. Softmax turns scores into probabilities
  2. Every position makes a guess
See Pro plansI have Pro: sign in
Next lessonAutoregressive generationPick a token, append it, run the model again. Start lesson 10

Transformers overviewSpotted a mistake? Tell us.