Attention and layers

In this chapter, we will learn what happens inside the layers. We will meet attention, the step in which every token collects information from the tokens before it. We will look at real attention in a real model, and then read the model's guess after each of its twelve layers to see where the answer first appears.

The problem: a word alone is not enough

The meaning of a token depends on the tokens around it. Take this sentence.

The animal didn't cross the street because it was too tired.

What is "it"? A reader knows at once: the animal. The embedding table of Chapter 5 cannot know. It gives "it" the same row in every sentence ever written. For the numbers of "it" to carry the animal, information must flow from one token to another.

Attention: every token looks back

Attention is the step in which each token takes in a weighted mix of itself and the tokens before it. The shares of the mix are worked out anew for every sentence. They always add up to 100%. These shares are not the weights of the model. The weights are fixed, and they are what works the shares out.

In words, the step works like this. Each token forms a question from its own numbers, such as "I am a pronoun, who is the noun I stand for?". Each earlier token, and the asking token too, offers a short description of itself. Where the question and a description fit well, that token gets a high share. Then the asking token adds to its own numbers a mix of what the others hold, each according to its share.

The fit between a question and a description is measured by a multiplication of two lists of numbers, the same kind as in Y = W · X. The scores are turned into shares with a formula we will meet in Chapter 7.

We looked at this step inside GPT-2 small.

(l) Attention: which earlier tokens does one token look at?
  "The animal didn't cross the street because it was too tired"
  looking from _it (token 9). It can only see the 9 tokens up to itself, never the ones after.
     layer 5, head 4:            The 8%  _animal 85%  _didn 2%  _cross 1%  _the 1%  _because 1%  _it 2%
     layer 1, head 1:            The 35%  _animal 6%  _didn 18%  't 7%  _cross 9%  _the 4%  _street 5%  _because 9%  _it 6%
     all 144 heads, averaged:    The 54%  _animal 5%  _didn 3%  't 4%  _cross 5%  _the 3%  _street 4%  _because 14%  _it 9%

Read the first line of numbers. In layer 5, one part of the model gives 85% of the attention of "it" to "animal". Nobody told it which word "it" stands for. It learned from text alone to look back at the noun the sentence is about. The program leaves out shares below 0.5%, which is why "'t" and "street" are missing from that line.

Heads: many questions at once

An attention head is one questioner. Each layer has several heads, and each head learned to ask a different question. GPT-2 small has 12 heads in each of its 12 layers.

heads in one layer     12
layers                 12
heads in the model     12 x 12  =  144

width of the model     768
numbers for one head   768 / 12  =  64

Each head works with a short list of 64 numbers for each token. The 12 results are put side by side, which gives 768 numbers again.

The second and third lines of the output show why we must be careful. Layer 1, head 1 spreads its look over everything. The average of all 144 heads looks mostly at the first word, which tells us little. One head that links "it" to "animal" does not prove that the model understands the sentence. It shows that one of 144 parts does something we can describe. For most heads nobody can say in a short sentence what they do.

Pick a token, then pick one of the four views: three single heads and the average. Guess first: which token does "street" look at hardest in layer 5, head 12?

Now go back to layer 5, head 4 and click through the tokens. "didn" and "cross" also give most of their attention to "animal". So this head is not a detector of pronouns. It looks back at the noun the sentence is about, from many places, and "it" is one of them. A short label for a head is always a little too neat.

A token never looks to the right

When the model works out the numbers of a token, it may only use that token and the ones before it. The reason is the loop of Chapter 1. While the model writes, the tokens to the right do not exist yet. The rule also holds for a text that is already complete, such as your prompt. The model was trained to guess the next token at every place of a text, and a look to the right would have shown it the answer.

This has a consequence that surprises people. Look at our sentence again. If it ended with "because it was too wide", then "it" would be the street. But when the model works on "it", the word "tired" or "wide" has not been read. The numbers of "it" are the same in both sentences. The model can only settle the question later, in the numbers of the tokens that come after.

The second half of a layer: thinking alone

After attention, each token is reworked alone, by two multiplications with a bend between them. This part is called the feed-forward step. It is the Y = W · X of Chapter 3 in its plainest form.

in           768 numbers of one token
W1           768 x 3,072           widen: 3,072 neurons each ask one question
bend         negative answers are pushed towards 0
W2         3,072 x 768             narrow again: back to 768 numbers
out          768 numbers, added to what the token had before

Attention moves information between tokens. The feed-forward step works on each token. Many researchers think of it as the place where the model keeps much of what it knows about the world. It holds 45.5% of all the weights of GPT-2 small. Attention holds 22.8%, which is half as much.

One layer, and twelve of them

A layer is one step of attention followed by one feed-forward step. GPT-2 small has 12 layers, one after the other.

          9 tokens x 768 numbers
                    |
        +-----------v-----------+
        |  attention            |   tokens collect from the tokens before them
        |  (12 heads)           |
        +-----------+-----------+
                    |  the result is ADDED to the numbers
        +-----------v-----------+
        |  feed-forward         |   each token is reworked alone
        +-----------+-----------+
                    |  the result is ADDED to the numbers
                    v
          9 tokens x 768 numbers        the same shape as before, other values

     ... and the same again, twelve times, each layer with weights of its own

A layer does not replace the numbers of a token. It adds a correction to them. The shape never changes on the way: 9 rows of 768 numbers go in, and 9 rows of 768 numbers come out.

(f) The shape of the numbers as they move through the model
  token ids                      9 numbers
  after the embedding lookup     9 x 768
  after layer 1                  9 x 768
  after layer 2                  9 x 768
  after layer 12                 9 x 768
  scores for the next token      50,257 numbers (one per token in the vocabulary)

This design is called a transformer. It was published in 2017, and nearly all the well known language models of today are transformers. The T in GPT stands for it. Our diagram leaves out one small step. Before each half of a layer the numbers are brought to a common scale, which is called layer normalisation. It holds 38,400 of the weights.

Watching the answer appear

The guess is built up step by step, from the first layer to the last. We return to the sentence of the book, "Maya lives in Paris, the capital of". After all 12 layers the model gives _France a chance of 28.81%. But when, on the way, did France come into play?

To find out, we used a trick of researchers. After every layer we took the numbers of the last token and asked: if the model had to guess now, what would it say? We asked by giving these numbers to the last step of the model. That step normally comes after layer 12 and turns 768 numbers into a score for every token.

(i) The guess takes shape layer by layer (the final read-out applied after each layer)
  after the embedding  _France   0.00 %    top 3: _destro 65%  _mathemat 26%  _challeng 8%
  after layer  1       _France   0.00 %    top 3: _the 63%  _course 10%  _a 2%
  after layer  2       _France   0.00 %    top 3: _the 82%  _course 7%  _a 1%
  after layer  3       _France   0.00 %    top 3: _the 75%  _course 14%  _a 1%
  after layer  4       _France   0.00 %    top 3: _the 82%  _course 2%  _Dam 1%
  after layer  5       _France   0.00 %    top 3: _the 88%  _a 1%  _his 1%
  after layer  6       _France   0.01 %    top 3: _the 55%  _town 5%  _British 3%
  after layer  7       _France   0.07 %    top 3: _the 36%  _South 6%  _Pennsylvania 5%
  after layer  8       _France   0.09 %    top 3: _the 66%  _North 6%  _South 3%
  after layer  9       _France  19.23 %    top 3: _the 39%  _France 19%  _Europe 8%
  after layer 10       _France  87.64 %    top 3: _France 88%  _Europe 6%  _the 2%
  after layer 11       _France  88.34 %    top 3: _France 88%  _the 5%  _Europe 2%
  after layer 12       _France  28.81 %    top 3: _France 29%  _the 25%  _a 5%

Move the slider through the layers yourself. It starts at layer 12, and 0 means before any layer. Guess first at which layer _France enters the top three.

The table tells a small story in four parts.

• Before any layer the guess is nonsense. The last token is "of", alone, and knows nothing of Paris.
• Layers 1 to 8 know grammar. After "of" comes a small word such as "the". From layer 6 on, words of place such as "town", "South" and "Pennsylvania" show up behind it, but France is still at 0.09%.
• Layers 9 and 10 bring the fact. The chance of France jumps from 0.09% to 19.23% and then to 87.64%. We also found the head that looks hardest from "of" back to "Paris". It sits in layer 9 and gives Paris 61% of its attention.
• Layer 12 pulls the confidence back to 28.81% and gives "the" its 25%. "The capital of the country" is also good English, and the last layer makes room for it.

Why this costs so much work

Every token must look at itself and at every token before it, in every head of every layer. That is why a long text costs much more work than a short one.

looks for a text of n tokens      1 + 2 + 3 + ... + n   =  n x (n + 1) / 2   in one head

a text of    9 tokens        9 x 10 / 2             =            45 looks
a text of 1,000 tokens   1,000 x 1,001 / 2          =       500,500 looks
in all 144 heads         500,500 x 144              =    72,072,000 looks

the text is 111 times longer   1,000 / 9            =  about 111
the looks are                  500,500 / 45         =  about 11,122 times as many

A text that is 111 times longer needs about 11,000 times as many looks. This growth is one of the main reasons why a model has a limit on how much text it can read at once. We return to that limit in Chapter 10.

python 03_gpt2_inside.py

32 x 16  =  512 heads

200 x 201 / 2  =  20,100 looks