Attention and transformers

Self-attention lets every token weigh every other token, so “it” can borrow meaning from “animal” or “street”. Change one word at the end of a sentence and watch the attention move.

10 min Free

Read this sentence: “The animal didn’t cross the street because it was too tired.” You know “it” is the animal. Change the last word to “wide” and “it” becomes the street. One word at the end changed the meaning of a word in the middle.

A model needs a way to do the same thing. When a token first becomes an embedding (see Embeddings), “it” has the same numbers in both sentences. Self-attention fixes that. For every token, it asks: how much does each other token affect how I should read this one? Then it mixes some of those tokens’ information into this token.

The transformer is the architecture built around this idea. It stacks many attention layers, and each layer has several attention heads that each learn their own pattern. Today’s LLMs, Gemini included, are built on transformers.

Press Play to see attention computed for “it”. Then switch Try it to wide, try the other head, and try Look back only.

  1. Tokens
  2. Query
  3. Scores
  4. Softmax
  5. Blend
  6. Every token
Tokens
Tokens

A sentence from Google's ML Crash Course. As a person you know what “it” means. Attention is how a transformer works that out. Press Play.

The scores are illustrative: hand-set to show what trained heads tend to learn. Masking, softmax and the weighted blend are the real maths. Real models use many heads in each of many layers.

What to notice

  1. Attention is a weighted average. Each token gets a score against every other token (query times key), softmax turns the scores into weights, and the token’s new value is a weighted mix. Nothing more mysterious than that, repeated many times.
  2. The weights depend on the whole sentence. Swap “tired” for “wide” and the weight moves from “animal” to “street”. The same word gets a different meaning from its context.
  3. Different heads learn different jobs. Google’s crash course gives examples: one head might resolve pronouns, another might track grammar. Head 2 here simply looks at the previous word. A model combines the output of all its heads.
  4. Decoders only look back. Models that generate text predict one token at a time, so each token may only attend to tokens before it. That’s the triangle you see with Look back only. Encoders, used for tasks like creating embeddings, can look both ways.
  5. Every token looks at every other token, so cost grows fast. Self-attention is roughly O(N²) in the number of tokens: double the prompt and that part of the work roughly quadruples. This is one reason long context costs more. See Context windows.

Layers build on layers
The output of one attention layer is the input to the next. According to Google’s crash course, earlier layers tend to pick up basic syntax, and deeper layers more abstract ideas such as sentiment. The single layer here is a small slice of that.

Check your understanding

1. In “…because it was too wide”, why does “it” attend most to “street”?
2. In a decoder that generates text, which tokens can “it” attend to?
3. Your prompt doubles from 4,000 to 8,000 tokens. Roughly what happens to the attention work in each layer?