How Transformers Changed AI: Attention Explained

In the sentence "The animal didn't cross the street because it was too tired", you know at once that "it" is the animal. A model has to work that out, and until 2017 the standard way of doing so was slow and forgetful. The attention mechanism replaced it, and every large language model in use today is built from it.

Attention has a reputation for being hard. The arithmetic is three lines. This guide does it by hand, then finds it at work inside a real model.

The problem it solved

Earlier models for language read a sentence one token at a time, carrying a running summary forward. Two things followed. Information from the start of a long passage faded by the end, and training could not be spread across many processors, because step ten had to wait for step nine.

The 2017 paper that introduced the Transformer proposed dropping that machinery altogether. Its abstract describes an architecture "based solely on attention mechanisms", and reports that it was more parallelizable and needed significantly less time to train (Vaswani et al.: Attention Is All You Need). That second property turned out to matter most. An architecture that trains in parallel can be trained on far more text.

What attention computes

Every token is represented by a vector, a list of numbers. Attention gives each token a new vector that is a blend of the others. The blend is decided by three things each token produces:

For one token, the calculation is:

1. Compare its query with every token's key, by multiplying and adding. A good match gives a high score.
2. Divide the scores by a constant, the square root of the vector size, to keep them in a workable range.
3. Turn the scores into weights that add up to 1.
4. Add up everyone's values, each multiplied by its weight.

That is scaled dot-product attention, the form used in the original paper (Harvard NLP: The Annotated Transformer). Nothing in it is specific to language. It is a weighted average in which the weights are worked out from the data.

By hand, then for real

The script does the four steps for three invented tokens. Then it loads GPT-2 and reads the attention weights the model actually used on the sentence about the animal.

import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Attention by hand: three tokens, each with a query, a key and a value of two numbers.
Q = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])    # what each token is looking for
K = np.array([[1.0, 0.0], [0.0, 1.0], [0.9, 0.1]])    # what each token offers
V = np.array([[10.0, 0.0], [0.0, 10.0], [5.0, 5.0]])  # what each token passes on
scores = Q @ K.T / np.sqrt(K.shape[1])
weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)   # softmax over each row
print("attention weights, each row sums to 1:")
print(weights.round(2))
print("new vector for each token, a weighted mix of the values:")
print((weights @ V).round(2))

# 2. The same operation inside a real model.
tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2", attn_implementation="eager").eval()
cfg = model.config
print()
print(f"GPT-2: {cfg.n_layer} layers x {cfg.n_head} heads = {cfg.n_layer * cfg.n_head} attention patterns, vectors of {cfg.n_embd} numbers")

ids = tok("The animal didn't cross the street because it was too tired", return_tensors="pt")
tokens = [tok.decode(i).strip() for i in ids.input_ids[0]]
with torch.no_grad():
    attention = model(**ids, output_attentions=True).attentions    # one tensor for each layer
it = tokens.index("it")
for layer, head in ((5, 4), (3, 3)):
    row = attention[layer - 1][0, head - 1, it, : it + 1]
    shown = "  ".join(f"{t} {w:.2f}" for t, w in zip(tokens, row.tolist()) if w >= 0.03)
    print(f"layer {layer}, head {head}, from 'it':  {shown}")

attention weights, each row sums to 1:
[[0.41 0.2  0.38]
 [0.24 0.49 0.26]
 [0.33 0.33 0.33]]
new vector for each token, a weighted mix of the values:
[[6.05 3.95]
 [3.75 6.25]
 [5.   5.  ]]

GPT-2: 12 layers x 12 heads = 144 attention patterns, vectors of 768 numbers
layer 5, head 4, from 'it':  The 0.08  animal 0.85
layer 3, head 3, from 'it':  The 0.04  the 0.06  street 0.09  because 0.71  it 0.08

• Each row of weights adds up to 1. The first token's query matches the first and third keys, so it takes 41% and 38% from them and only 20% from the second.
• The third token wants both things equally, so it gets a third of each, and its new vector is an even blend.
• GPT-2 has 144 of these patterns, 12 heads in each of 12 layers, all running on the same sentence.
• Layer 5, head 4 resolves the pronoun. From "it", 85% of the weight lands on "animal".
• Layer 3, head 3 does a different job. From "it", 71% lands on "because", the word just before.

Nobody wrote a rule that says pronouns should look for nouns. The queries, keys and values are produced by parameters that were adjusted during training, and a head that links pronouns to nouns emerged because it helped predict the next word.

Heads, layers and order

Heads. One attention pattern can track one kind of relationship. Running several side by side lets the model follow several at once, such as grammar, reference and position. This is multi-head attention.

Layers. A layer is attention followed by a small ordinary network applied to each token, with shortcuts that let information skip past. Layers are stacked, so later ones work on the output of earlier ones.

Order. Attention on its own does not know which token came first. Shuffle the sentence and each token would end up with exactly the same new vector as before. So information about position is added to each token's vector before the first layer.

Three families

The original Transformer had two halves, an encoder that reads and a decoder that writes. Later models kept one half or both (Hugging Face: how do Transformers work):

BERT reads in both directions at once (Devlin et al.: BERT). GPT-style models are restricted to what came before, because their job is to predict what comes next. That restriction is visible in the output above: the weights from "it" cover only the words up to "it".

The chat models people use are decoder-only. See how LLMs work for what such a model does with the result. Encoder models are behind embeddings.

What it costs

Look at the first table of numbers in the output: three tokens gave a three-by-three grid of weights. Every token is compared with every other. A thousand tokens means a million comparisons in each head, and ten thousand tokens means a hundred million.

This is the main reason long inputs are slower and cost more, and one reason models have a limit on how much text they accept. See context windows explained.

A caution about reading attention

It is tempting to treat attention weights as the model's reasoning made visible. They are less than that. A study of this question found that attention weights often do not agree with other measures of which inputs mattered, and that quite different weights can produce the same prediction (Jain and Wallace: Attention is not Explanation).

The head shown here does point from "it" to "animal". That is evidence about what one part of the model does. It is not a full account of why the model would continue the sentence one way or another.

Common misunderstandings

• "Attention is like human attention." It is a weighted average. The name is a metaphor.
• "A transformer reads left to right." In training it processes every position at once. Only generation is one token at a time.
• "Word order comes for free." Position has to be added explicitly.
• "The model stores a rule for pronouns." It stores numbers. The behaviour is learned, spread over many heads, and not guaranteed.
• "Attention explains the answer." It shows where a head looked.

The Transformer architecture, with self and cross attention, is covered in depth in the Generative AI course.

Sources

• Attention Is All You Need (arXiv) — Vaswani et al.
• The Annotated Transformer (Harvard NLP) — Harvard NLP
• How do Transformers work? (Hugging Face LLM Course) — Hugging Face
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (arXiv) — Devlin et al.
• Attention is not Explanation (arXiv) — Jain and Wallace